from __future__ import division, print_function

import unittest
import numpy as np

from keras import backend as K

from rvseg import loss

class TestModel(unittest.TestCase):
    def sample_tensors(self):
        # have shapes (height, width, classes=2)
        y_true = K.constant([[[0, 1], [1, 0]],
                             [[1, 0], [1, 0]]])
        y_pred = K.constant([[[.4, .6], [.1, .9]],
                             [[.5, .5], [.0, 1.]]])
        return y_true, y_pred

    def test_sorensen_dice(self):
        y_true, y_pred = self.sample_tensors()
        dice_coefs = loss.soft_sorensen_dice(y_true, y_pred, axis=[0, 1])
        dice_coefs = K.eval(dice_coefs)
        # class 1: (2 * 0.6 + 1) / (1 + 3 + 1) = 0.44
        # class 2: (2 * 0.6 + 1) / (3 + 1 + 1) = 0.44
        expected_dice_coefs = [0.44, 0.44]
        for x,y in zip(dice_coefs, expected_dice_coefs):
            self.assertAlmostEqual(x, y)

        dice_coefs = loss.hard_sorensen_dice(y_true, y_pred, axis=[0, 1])
        dice_coefs = K.eval(dice_coefs)
        # class 1: (2 * 0 + 1) / (0 + 3 + 1) = 0.25
        # class 2: (2 * 1 + 1) / (3 + 1 + 1) = 0.6
        expected_dice_coefs = [0.25, 0.6]
        for x,y in zip(dice_coefs, expected_dice_coefs):
            self.assertAlmostEqual(x, y)

    def test_weighted_categorical_crossentropy(self):
        weights = [1, 9]
        y_true, y_pred = self.sample_tensors()

        lossfunc = loss.weighted_categorical_crossentropy(
            y_true, y_pred, weights)
        loss_val = K.eval(lossfunc)

        w = 2 * np.array(weights) / sum(weights)
        logs = -np.log(np.array([.1, .4, .5, .6, .9, 1e-8]))
        expected_loss_val = w[1]*logs[3]/4 + w[0]*(logs[0] + logs[2] + logs[5])/4

        self.assertAlmostEqual(np.mean(expected_loss_val), loss_val, places=5)