# Copyright 2020 The FastEstimator Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
import unittest

import numpy as np
import tensorflow as tf
import torch

import fastestimator as fe


class TestBinarayCrossEntropy(unittest.TestCase):
    @classmethod
    def setUpClass(cls):
        cls.tf_true = tf.constant([[1], [0], [1], [0]])
        cls.tf_pred = tf.constant([[0.9], [0.3], [0.8], [0.1]])
        cls.torch_true = torch.tensor([[1], [0], [1], [0]])
        cls.torch_pred = torch.tensor([[0.9], [0.3], [0.8], [0.1]])

    def test_binaray_crossentropy_average_loss_true_tf(self):
        obj1 = fe.backend.binary_crossentropy(y_pred=self.tf_pred, y_true=self.tf_true).numpy()
        obj2 = 0.19763474
        self.assertTrue(np.allclose(obj1, obj2))

    def test_binaray_crossentropy_average_loss_false_tf(self):
        obj1 = fe.backend.binary_crossentropy(y_pred=self.tf_pred, y_true=self.tf_true, average_loss=False).numpy()
        obj2 = np.array([0.10536041, 0.3566748, 0.22314338, 0.10536041])
        self.assertTrue(np.allclose(obj1, obj2))

    def test_binaray_crossentropy_average_from_logit_average_loss_true_tf(self):
        obj1 = fe.backend.binary_crossentropy(y_pred=self.tf_pred,
                                              y_true=self.tf_true,
                                              from_logits=True,
                                              average_loss=True).numpy()
        obj2 = 0.57775164
        self.assertTrue(np.allclose(obj1, obj2))

    def test_binaray_crossentropy_average_from_logit_average_loss_false_tf(self):
        obj1 = fe.backend.binary_crossentropy(y_pred=self.tf_pred,
                                              y_true=self.tf_true,
                                              from_logits=True,
                                              average_loss=False).numpy()
        obj2 = np.array([0.34115386, 0.8543553, 0.37110066, 0.7443967])
        self.assertTrue(np.allclose(obj1, obj2))

    def test_binaray_crossentropy_average_loss_true_torch(self):
        obj1 = fe.backend.binary_crossentropy(y_pred=self.torch_pred, y_true=self.torch_true).numpy()
        obj2 = 0.19763474
        self.assertTrue(np.allclose(obj1, obj2))

    def test_binaray_crossentropy_average_loss_false_torch(self):
        obj1 = fe.backend.binary_crossentropy(y_pred=self.torch_pred, y_true=self.torch_true,
                                              average_loss=False).numpy()
        obj2 = np.array([0.10536041, 0.3566748, 0.22314338, 0.10536041])
        self.assertTrue(np.allclose(obj1, obj2))

    def test_binaray_crossentropy_average_from_logit_average_loss_true_torch(self):
        obj1 = fe.backend.binary_crossentropy(y_pred=self.torch_pred,
                                              y_true=self.torch_true,
                                              from_logits=True,
                                              average_loss=True).numpy()
        obj2 = 0.57775164
        self.assertTrue(np.allclose(obj1, obj2))

    def test_binaray_crossentropy_average_from_logit_average_loss_false_torch(self):
        obj1 = fe.backend.binary_crossentropy(y_pred=self.torch_pred,
                                              y_true=self.torch_true,
                                              from_logits=True,
                                              average_loss=False).numpy()
        obj2 = np.array([0.34115386, 0.8543553, 0.37110066, 0.7443967])
        self.assertTrue(np.allclose(obj1, obj2))

[((909, 942), 'tensorflow.constant', 'tf.constant', (['[[1], [0], [1], [0]]'], {}), '([[1], [0], [1], [0]])\n', (920, 942), True, 'import tensorflow as tf\n'), ((965, 1006), 'tensorflow.constant', 'tf.constant', (['[[0.9], [0.3], [0.8], [0.1]]'], {}), '([[0.9], [0.3], [0.8], [0.1]])\n', (976, 1006), True, 'import tensorflow as tf\n'), ((1032, 1066), 'torch.tensor', 'torch.tensor', (['[[1], [0], [1], [0]]'], {}), '([[1], [0], [1], [0]])\n', (1044, 1066), False, 'import torch\n'), ((1092, 1134), 'torch.tensor', 'torch.tensor', (['[[0.9], [0.3], [0.8], [0.1]]'], {}), '([[0.9], [0.3], [0.8], [0.1]])\n', (1104, 1134), False, 'import torch\n'), ((1564, 1621), 'numpy.array', 'np.array', (['[0.10536041, 0.3566748, 0.22314338, 0.10536041]'], {}), '([0.10536041, 0.3566748, 0.22314338, 0.10536041])\n', (1572, 1621), True, 'import numpy as np\n'), ((2469, 2525), 'numpy.array', 'np.array', (['[0.34115386, 0.8543553, 0.37110066, 0.7443967]'], {}), '([0.34115386, 0.8543553, 0.37110066, 0.7443967])\n', (2477, 2525), True, 'import numpy as np\n'), ((3068, 3125), 'numpy.array', 'np.array', (['[0.10536041, 0.3566748, 0.22314338, 0.10536041]'], {}), '([0.10536041, 0.3566748, 0.22314338, 0.10536041])\n', (3076, 3125), True, 'import numpy as np\n'), ((3991, 4047), 'numpy.array', 'np.array', (['[0.34115386, 0.8543553, 0.37110066, 0.7443967]'], {}), '([0.34115386, 0.8543553, 0.37110066, 0.7443967])\n', (3999, 4047), True, 'import numpy as np\n'), ((1344, 1367), 'numpy.allclose', 'np.allclose', (['obj1', 'obj2'], {}), '(obj1, obj2)\n', (1355, 1367), True, 'import numpy as np\n'), ((1646, 1669), 'numpy.allclose', 'np.allclose', (['obj1', 'obj2'], {}), '(obj1, obj2)\n', (1657, 1669), True, 'import numpy as np\n'), ((2074, 2097), 'numpy.allclose', 'np.allclose', (['obj1', 'obj2'], {}), '(obj1, obj2)\n', (2085, 2097), True, 'import numpy as np\n'), ((2550, 2573), 'numpy.allclose', 'np.allclose', (['obj1', 'obj2'], {}), '(obj1, obj2)\n', (2561, 2573), True, 'import numpy as np\n'), ((2793, 2816), 'numpy.allclose', 'np.allclose', (['obj1', 'obj2'], {}), '(obj1, obj2)\n', (2804, 2816), True, 'import numpy as np\n'), ((3150, 3173), 'numpy.allclose', 'np.allclose', (['obj1', 'obj2'], {}), '(obj1, obj2)\n', (3161, 3173), True, 'import numpy as np\n'), ((3587, 3610), 'numpy.allclose', 'np.allclose', (['obj1', 'obj2'], {}), '(obj1, obj2)\n', (3598, 3610), True, 'import numpy as np\n'), ((4072, 4095), 'numpy.allclose', 'np.allclose', (['obj1', 'obj2'], {}), '(obj1, obj2)\n', (4083, 4095), True, 'import numpy as np\n'), ((1213, 1285), 'fastestimator.backend.binary_crossentropy', 'fe.backend.binary_crossentropy', ([], {'y_pred': 'self.tf_pred', 'y_true': 'self.tf_true'}), '(y_pred=self.tf_pred, y_true=self.tf_true)\n', (1243, 1285), True, 'import fastestimator as fe\n'), ((1448, 1544), 'fastestimator.backend.binary_crossentropy', 'fe.backend.binary_crossentropy', ([], {'y_pred': 'self.tf_pred', 'y_true': 'self.tf_true', 'average_loss': '(False)'}), '(y_pred=self.tf_pred, y_true=self.tf_true,\n    average_loss=False)\n', (1478, 1544), True, 'import fastestimator as fe\n'), ((1768, 1881), 'fastestimator.backend.binary_crossentropy', 'fe.backend.binary_crossentropy', ([], {'y_pred': 'self.tf_pred', 'y_true': 'self.tf_true', 'from_logits': '(True)', 'average_loss': '(True)'}), '(y_pred=self.tf_pred, y_true=self.tf_true,\n    from_logits=True, average_loss=True)\n', (1798, 1881), True, 'import fastestimator as fe\n'), ((2197, 2311), 'fastestimator.backend.binary_crossentropy', 'fe.backend.binary_crossentropy', ([], {'y_pred': 'self.tf_pred', 'y_true': 'self.tf_true', 'from_logits': '(True)', 'average_loss': '(False)'}), '(y_pred=self.tf_pred, y_true=self.tf_true,\n    from_logits=True, average_loss=False)\n', (2227, 2311), True, 'import fastestimator as fe\n'), ((2656, 2734), 'fastestimator.backend.binary_crossentropy', 'fe.backend.binary_crossentropy', ([], {'y_pred': 'self.torch_pred', 'y_true': 'self.torch_true'}), '(y_pred=self.torch_pred, y_true=self.torch_true)\n', (2686, 2734), True, 'import fastestimator as fe\n'), ((2900, 3003), 'fastestimator.backend.binary_crossentropy', 'fe.backend.binary_crossentropy', ([], {'y_pred': 'self.torch_pred', 'y_true': 'self.torch_true', 'average_loss': '(False)'}), '(y_pred=self.torch_pred, y_true=self.\n    torch_true, average_loss=False)\n', (2930, 3003), True, 'import fastestimator as fe\n'), ((3275, 3395), 'fastestimator.backend.binary_crossentropy', 'fe.backend.binary_crossentropy', ([], {'y_pred': 'self.torch_pred', 'y_true': 'self.torch_true', 'from_logits': '(True)', 'average_loss': '(True)'}), '(y_pred=self.torch_pred, y_true=self.\n    torch_true, from_logits=True, average_loss=True)\n', (3305, 3395), True, 'import fastestimator as fe\n'), ((3713, 3834), 'fastestimator.backend.binary_crossentropy', 'fe.backend.binary_crossentropy', ([], {'y_pred': 'self.torch_pred', 'y_true': 'self.torch_true', 'from_logits': '(True)', 'average_loss': '(False)'}), '(y_pred=self.torch_pred, y_true=self.\n    torch_true, from_logits=True, average_loss=False)\n', (3743, 3834), True, 'import fastestimator as fe\n')]