# 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))

[((27, 22, 27, 55), 'tensorflow.constant', 'tf.constant', ({(27, 34, 27, 54): '[[1], [0], [1], [0]]'}, {}), '([[1], [0], [1], [0]])', True, 'import tensorflow as tf\n'), ((28, 22, 28, 63), 'tensorflow.constant', 'tf.constant', ({(28, 34, 28, 62): '[[0.9], [0.3], [0.8], [0.1]]'}, {}), '([[0.9], [0.3], [0.8], [0.1]])', True, 'import tensorflow as tf\n'), ((29, 25, 29, 59), 'torch.tensor', 'torch.tensor', ({(29, 38, 29, 58): '[[1], [0], [1], [0]]'}, {}), '([[1], [0], [1], [0]])', False, 'import torch\n'), ((30, 25, 30, 67), 'torch.tensor', 'torch.tensor', ({(30, 38, 30, 66): '[[0.9], [0.3], [0.8], [0.1]]'}, {}), '([[0.9], [0.3], [0.8], [0.1]])', False, 'import torch\n'), ((39, 15, 39, 72), 'numpy.array', 'np.array', ({(39, 24, 39, 71): '[0.10536041, 0.3566748, 0.22314338, 0.10536041]'}, {}), '([0.10536041, 0.3566748, 0.22314338, 0.10536041])', True, 'import numpy as np\n'), ((55, 15, 55, 71), 'numpy.array', 'np.array', ({(55, 24, 55, 70): '[0.34115386, 0.8543553, 0.37110066, 0.7443967]'}, {}), '([0.34115386, 0.8543553, 0.37110066, 0.7443967])', True, 'import numpy as np\n'), ((66, 15, 66, 72), 'numpy.array', 'np.array', ({(66, 24, 66, 71): '[0.10536041, 0.3566748, 0.22314338, 0.10536041]'}, {}), '([0.10536041, 0.3566748, 0.22314338, 0.10536041])', True, 'import numpy as np\n'), ((82, 15, 82, 71), 'numpy.array', 'np.array', ({(82, 24, 82, 70): '[0.34115386, 0.8543553, 0.37110066, 0.7443967]'}, {}), '([0.34115386, 0.8543553, 0.37110066, 0.7443967])', True, 'import numpy as np\n'), ((35, 24, 35, 47), 'numpy.allclose', 'np.allclose', ({(35, 36, 35, 40): 'obj1', (35, 42, 35, 46): 'obj2'}, {}), '(obj1, obj2)', True, 'import numpy as np\n'), ((40, 24, 40, 47), 'numpy.allclose', 'np.allclose', ({(40, 36, 40, 40): 'obj1', (40, 42, 40, 46): 'obj2'}, {}), '(obj1, obj2)', True, 'import numpy as np\n'), ((48, 24, 48, 47), 'numpy.allclose', 'np.allclose', ({(48, 36, 48, 40): 'obj1', (48, 42, 48, 46): 'obj2'}, {}), '(obj1, obj2)', True, 'import numpy as np\n'), ((56, 24, 56, 47), 'numpy.allclose', 'np.allclose', ({(56, 36, 56, 40): 'obj1', (56, 42, 56, 46): 'obj2'}, {}), '(obj1, obj2)', True, 'import numpy as np\n'), ((61, 24, 61, 47), 'numpy.allclose', 'np.allclose', ({(61, 36, 61, 40): 'obj1', (61, 42, 61, 46): 'obj2'}, {}), '(obj1, obj2)', True, 'import numpy as np\n'), ((67, 24, 67, 47), 'numpy.allclose', 'np.allclose', ({(67, 36, 67, 40): 'obj1', (67, 42, 67, 46): 'obj2'}, {}), '(obj1, obj2)', True, 'import numpy as np\n'), ((75, 24, 75, 47), 'numpy.allclose', 'np.allclose', ({(75, 36, 75, 40): 'obj1', (75, 42, 75, 46): 'obj2'}, {}), '(obj1, obj2)', True, 'import numpy as np\n'), ((83, 24, 83, 47), 'numpy.allclose', 'np.allclose', ({(83, 36, 83, 40): 'obj1', (83, 42, 83, 46): 'obj2'}, {}), '(obj1, obj2)', True, 'import numpy as np\n'), ((33, 15, 33, 87), 'fastestimator.backend.binary_crossentropy', 'fe.backend.binary_crossentropy', (), '', True, 'import fastestimator as fe\n'), ((38, 15, 38, 107), 'fastestimator.backend.binary_crossentropy', 'fe.backend.binary_crossentropy', (), '', True, 'import fastestimator as fe\n'), ((43, 15, 46, 64), 'fastestimator.backend.binary_crossentropy', 'fe.backend.binary_crossentropy', (), '', True, 'import fastestimator as fe\n'), ((51, 15, 54, 65), 'fastestimator.backend.binary_crossentropy', 'fe.backend.binary_crossentropy', (), '', True, 'import fastestimator as fe\n'), ((59, 15, 59, 93), 'fastestimator.backend.binary_crossentropy', 'fe.backend.binary_crossentropy', (), '', True, 'import fastestimator as fe\n'), ((64, 15, 65, 65), 'fastestimator.backend.binary_crossentropy', 'fe.backend.binary_crossentropy', (), '', True, 'import fastestimator as fe\n'), ((70, 15, 73, 64), 'fastestimator.backend.binary_crossentropy', 'fe.backend.binary_crossentropy', (), '', True, 'import fastestimator as fe\n'), ((78, 15, 81, 65), 'fastestimator.backend.binary_crossentropy', 'fe.backend.binary_crossentropy', (), '', True, 'import fastestimator as fe\n')]