# pylint: disable=missing-docstring
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
from __future__ import unicode_literals

import unittest
import numpy as np

# Weird imports / assignment because the normal import syntax doesn't work for tf.keras in tf 1.8
from tensorflow import keras

# pylint:disable=wrong-import-position
Sequential = keras.models.Sequential
Dense = keras.layers.Dense
Activation = keras.layers.Activation

from cleverhans.utils_keras import KerasModelWrapper


class TestKerasModelWrapper(unittest.TestCase):
    def setUp(self):
        import tensorflow as tf

        def dummy_model():
            input_shape = (100,)
            return Sequential(
                [
                    Dense(20, name="l1", input_shape=input_shape),
                    Dense(10, name="l2"),
                    Activation("softmax", name="softmax"),
                ]
            )

        self.sess = tf.Session()
        self.sess.as_default()
        self.model = dummy_model()

    def test_softmax_layer_name_is_softmax(self):
        model = KerasModelWrapper(self.model)
        softmax_name = model._get_softmax_name()
        self.assertEqual(softmax_name, "softmax")

    def test_logit_layer_name_is_logits(self):
        model = KerasModelWrapper(self.model)
        logits_name = model._get_logits_name()
        self.assertEqual(logits_name, "l2")

    def test_get_logits(self):
        import tensorflow as tf

        model = KerasModelWrapper(self.model)
        x = tf.placeholder(tf.float32, shape=(None, 100))
        preds = model.get_probs(x)
        logits = model.get_logits(x)

        x_val = np.random.rand(2, 100)
        tf.global_variables_initializer().run(session=self.sess)
        p_val, logits = self.sess.run([preds, logits], feed_dict={x: x_val})
        p_gt = np.exp(logits) / np.sum(np.exp(logits), axis=1, keepdims=True)
        self.assertTrue(np.allclose(p_val, p_gt, atol=1e-6))

    def test_get_probs(self):
        import tensorflow as tf

        model = KerasModelWrapper(self.model)
        x = tf.placeholder(tf.float32, shape=(None, 100))
        preds = model.get_probs(x)

        x_val = np.random.rand(2, 100)
        tf.global_variables_initializer().run(session=self.sess)
        p_val = self.sess.run(preds, feed_dict={x: x_val})
        self.assertTrue(np.allclose(np.sum(p_val, axis=1), 1, atol=1e-6))
        self.assertTrue(np.all(p_val >= 0))
        self.assertTrue(np.all(p_val <= 1))

    def test_get_layer_names(self):
        model = KerasModelWrapper(self.model)
        layer_names = model.get_layer_names()
        self.assertEqual(layer_names, ["l1", "l2", "softmax"])

    def test_fprop(self):
        import tensorflow as tf

        model = KerasModelWrapper(self.model)
        x = tf.placeholder(tf.float32, shape=(None, 100))
        out_dict = model.fprop(x)

        self.assertEqual(set(out_dict.keys()), set(["l1", "l2", "softmax"]))
        # Test the dimension of the hidden represetation
        self.assertEqual(int(out_dict["l1"].shape[1]), 20)
        self.assertEqual(int(out_dict["l2"].shape[1]), 10)

        # Test the caching
        x2 = tf.placeholder(tf.float32, shape=(None, 100))
        out_dict2 = model.fprop(x2)
        self.assertEqual(set(out_dict2.keys()), set(["l1", "l2", "softmax"]))
        self.assertEqual(int(out_dict2["l1"].shape[1]), 20)


if __name__ == "__main__":
    unittest.main()
