from unittest import TestCase

import numpy as np
import keras.backend as K


from fetal_net.metrics import weighted_dice_coefficient


class TestWeightedDice(TestCase):
    def test_weighted_dice_coefficient(self):
        data = np.zeros((5**3) * 3).reshape(3, 5, 5, 5)
        data[0, 0:1] = 1
        data[1, 0:2] = 1
        data[2, 1:4] = 1

        max_dice = K.eval(weighted_dice_coefficient(K.variable(data), K.variable(data)))
        for index in range(data.shape[0]):
            temp_data = np.copy(data)
            temp_data[index] = 0
            dice = K.eval(weighted_dice_coefficient(K.variable(data), K.variable(temp_data)))
            self.assertAlmostEqual(dice, (2 * max_dice)/3, delta=0.00001)

    def test_blank_dice_coefficient(self):
        data = np.zeros((5**3) * 3).reshape(3, 5, 5, 5)
        blank = np.copy(data)
        data[0, 0:1] = 1
        data[1, 0:2] = 1
        data[2, 1:4] = 1

        self.assertAlmostEqual(K.eval(weighted_dice_coefficient(K.variable(data), K.variable(blank))), 0, delta=0.00001)

    def test_empty_label(self):
        data = np.zeros((5**3) * 3).reshape(3, 5, 5, 5)
        data[1, 0:2] = 1
        data[2, 1:4] = 1

        self.assertEqual(K.eval(weighted_dice_coefficient(K.variable(data), K.variable(data))), 1)