# 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, 'tensorflow.constant', 'tf.constant', (['[[1], [0], [1], [0]]'], {}), True, 'import tensorflow as tf\n'), (28, 'tensorflow.constant', 'tf.constant', (['[[0.9], [0.3], [0.8], [0.1]]'], {}), True, 'import tensorflow as tf\n'), (29, 'torch.tensor', 'torch.tensor', (['[[1], [0], [1], [0]]'], {}), False, 'import torch\n'), (30, 'torch.tensor', 'torch.tensor', (['[[0.9], [0.3], [0.8], [0.1]]'], {}), False, 'import torch\n'), (39, 'numpy.array', 'np.array', (['[0.10536041, 0.3566748, 0.22314338, 0.10536041]'], {}), True, 'import numpy as np\n'), (55, 'numpy.array', 'np.array', (['[0.34115386, 0.8543553, 0.37110066, 0.7443967]'], {}), True, 'import numpy as np\n'), (66, 'numpy.array', 'np.array', (['[0.10536041, 0.3566748, 0.22314338, 0.10536041]'], {}), True, 'import numpy as np\n'), (82, 'numpy.array', 'np.array', (['[0.34115386, 0.8543553, 0.37110066, 0.7443967]'], {}), True, 'import numpy as np\n'), (35, 'numpy.allclose', 'np.allclose', (['obj1', 'obj2'], {}), True, 'import numpy as np\n'), (40, 'numpy.allclose', 'np.allclose', (['obj1', 'obj2'], {}), True, 'import numpy as np\n'), (48, 'numpy.allclose', 'np.allclose', (['obj1', 'obj2'], {}), True, 'import numpy as np\n'), (56, 'numpy.allclose', 'np.allclose', (['obj1', 'obj2'], {}), True, 'import numpy as np\n'), (61, 'numpy.allclose', 'np.allclose', (['obj1', 'obj2'], {}), True, 'import numpy as np\n'), (67, 'numpy.allclose', 'np.allclose', (['obj1', 'obj2'], {}), True, 'import numpy as np\n'), (75, 'numpy.allclose', 'np.allclose', (['obj1', 'obj2'], {}), True, 'import numpy as np\n'), (83, 'numpy.allclose', 'np.allclose', (['obj1', 'obj2'], {}), True, 'import numpy as np\n'), (33, 'fastestimator.backend.binary_crossentropy', 'fe.backend.binary_crossentropy', ([], {'y_pred': 'self.tf_pred', 'y_true': 'self.tf_true'}), True, 'import fastestimator as fe\n'), (38, 'fastestimator.backend.binary_crossentropy', 'fe.backend.binary_crossentropy', ([], {'y_pred': 'self.tf_pred', 'y_true': 'self.tf_true', 'average_loss': '(False)'}), True, 'import fastestimator as fe\n'), (43, 'fastestimator.backend.binary_crossentropy', 'fe.backend.binary_crossentropy', ([], {'y_pred': 'self.tf_pred', 'y_true': 'self.tf_true', 'from_logits': '(True)', 'average_loss': '(True)'}), True, 'import fastestimator as fe\n'), (51, 'fastestimator.backend.binary_crossentropy', 'fe.backend.binary_crossentropy', ([], {'y_pred': 'self.tf_pred', 'y_true': 'self.tf_true', 'from_logits': '(True)', 'average_loss': '(False)'}), True, 'import fastestimator as fe\n'), (59, 'fastestimator.backend.binary_crossentropy', 'fe.backend.binary_crossentropy', ([], {'y_pred': 'self.torch_pred', 'y_true': 'self.torch_true'}), True, 'import fastestimator as fe\n'), (64, 'fastestimator.backend.binary_crossentropy', 'fe.backend.binary_crossentropy', ([], {'y_pred': 'self.torch_pred', 'y_true': 'self.torch_true', 'average_loss': '(False)'}), True, 'import fastestimator as fe\n'), (70, 'fastestimator.backend.binary_crossentropy', 'fe.backend.binary_crossentropy', ([], {'y_pred': 'self.torch_pred', 'y_true': 'self.torch_true', 'from_logits': '(True)', 'average_loss': '(True)'}), True, 'import fastestimator as fe\n'), (78, 'fastestimator.backend.binary_crossentropy', 'fe.backend.binary_crossentropy', ([], {'y_pred': 'self.torch_pred', 'y_true': 'self.torch_true', 'from_logits': '(True)', 'average_loss': '(False)'}), True, 'import fastestimator as fe\n')]