import unittest
import torch
import random
from math import isclose
from unittest.mock import Mock, call

from tests.test_transforms import chech_data_preservation
from rising.transforms.intensity import *


class MyTestCase(unittest.TestCase):
    def setUp(self) -> None:
        torch.manual_seed(0)
        random.seed(0)
        self.batch_dict = {
            "data": torch.arange(1, 10).reshape(1, 1, 3, 3).float(),
            "seg": torch.rand(1, 1, 3, 3),
            "label": torch.arange(3)
        }

    def test_clamp_transform(self):
        trafo = ClampTransform(0, 1)

        self.assertTrue(chech_data_preservation(trafo, self.batch_dict))

        outp = trafo(**self.batch_dict)
        self.assertTrue((outp["data"] == torch.ones_like(outp["data"])).all())

    def test_norm_range_transform(self):
        trafo = NormRangeTransform(0.1, 0.2, per_channel=False)
        self.assertTrue(chech_data_preservation(trafo, self.batch_dict))

        trafo = NormRangeTransform(0.1, 0.2, per_channel=True)
        self.assertTrue(chech_data_preservation(trafo, self.batch_dict))

        outp = trafo(**self.batch_dict)
        self.assertTrue(isclose(outp["data"].min().item(), 0.1, abs_tol=1e-6))
        self.assertTrue(isclose(outp["data"].max().item(), 0.2, abs_tol=1e-6))

    def test_norm_min_max_transform(self):
        trafo = NormMinMaxTransform(per_channel=False)
        self.assertTrue(chech_data_preservation(trafo, self.batch_dict))

        trafo = NormMinMaxTransform(per_channel=True)
        self.assertTrue(chech_data_preservation(trafo, self.batch_dict))

        outp = trafo(**self.batch_dict)
        self.assertTrue(isclose(outp["data"].min().item(), 0., abs_tol=1e-6))
        self.assertTrue(isclose(outp["data"].max().item(), 1., abs_tol=1e-6))

    def test_norm_zero_mean_transform(self):
        trafo = NormZeroMeanUnitStdTransform(per_channel=False)
        self.assertTrue(chech_data_preservation(trafo, self.batch_dict))

        trafo = NormZeroMeanUnitStdTransform(per_channel=True)
        self.assertTrue(chech_data_preservation(trafo, self.batch_dict))

        outp = trafo(**self.batch_dict)
        self.assertTrue(isclose(outp["data"].mean().item(), 0., abs_tol=1e-6))
        self.assertTrue(isclose(outp["data"].std().item(), 1., abs_tol=1e-6))

    def test_norm_std_transform(self):
        mean = self.batch_dict["data"].mean().item()
        std = self.batch_dict["data"].std().item()
        trafo = NormMeanStdTransform(mean=mean, std=std, per_channel=False)
        self.assertTrue(chech_data_preservation(trafo, self.batch_dict))

        trafo = NormMeanStdTransform(mean=mean, std=std, per_channel=True)
        self.assertTrue(chech_data_preservation(trafo, self.batch_dict))

        outp = trafo(**self.batch_dict)
        self.assertTrue(isclose(outp["data"].mean().item(), 0., abs_tol=1e-6))
        self.assertTrue(isclose(outp["data"].std().item(), 1., abs_tol=1e-6))

    def test_noise_transform(self):
        trafo = NoiseTransform('normal', mean=75, std=1)
        self.assertTrue(chech_data_preservation(trafo, self.batch_dict))
        self.check_noise_distance(trafo)

    def test_expoential_noise_transform(self):
        trafo = ExponentialNoiseTransform(lambd=0.0001)
        self.assertTrue(chech_data_preservation(trafo, self.batch_dict))
        self.check_noise_distance(trafo)

    def test_gaussian_noise_transform(self):
        trafo = GaussianNoiseTransform(mean=75, std=1)
        self.assertTrue(chech_data_preservation(trafo, self.batch_dict))
        self.check_noise_distance(trafo)

    def check_noise_distance(self, trafo, min_diff=50):
        outp = trafo(**self.batch_dict)
        comp_diff = (outp["data"] - self.batch_dict["data"]).mean().item()
        self.assertTrue(comp_diff > min_diff)

    def test_per_channel_transform_per_channel_true(self):
        mock = Mock(return_value=0)

        def augment_fn(inp, *args, **kwargs):
            return mock(inp)

        trafo = RandomValuePerChannelTransform(
            augment_fn, random_mode="random", per_channel=True, keys=('label',))
        self.batch_dict["label"] = self.batch_dict["label"][None]
        output = trafo(**self.batch_dict)
        calls = [call(torch.tensor([0])), call(torch.tensor([1])),
                 call(torch.tensor([2])), ]
        mock.assert_has_calls(calls)

    def test_random_add_value(self):
        trafo = RandomAddValue("random")
        self.assertTrue(chech_data_preservation(trafo, self.batch_dict))

        random.seed(0)
        rand_val = random.random()
        random.seed(0)
        outp = trafo(**self.batch_dict)
        expected_out = self.batch_dict["data"] + rand_val
        self.assertTrue((outp["data"] == expected_out).all())
        self.assertEqual(trafo.random_mode, "random")

    def test_random_scale_value(self):
        trafo = RandomScaleValue("random")
        self.assertTrue(chech_data_preservation(trafo, self.batch_dict))

        random.seed(0)
        rand_val = random.random()
        random.seed(0)
        outp = trafo(**self.batch_dict)
        expected_out = self.batch_dict["data"] * rand_val
        self.assertTrue((outp["data"] == expected_out).all())
        self.assertEqual(trafo.random_mode, "random")

    def test_gamma_transform_scalar(self):
        trafo = GammaCorrectionTransform(gamma=2)
        self.assertTrue(chech_data_preservation(trafo, self.batch_dict))

        trafo = GammaCorrectionTransform(gamma=2)
        outp = trafo(**self.batch_dict)
        expected_out = self.batch_dict["data"].pow(2)
        self.assertTrue((outp["data"] == expected_out).all())

    def test_gamma_transform_error(self):
        with self.assertRaises(TypeError):
            trafo = GammaCorrectionTransform(gamma=(1, 1, 1))

    def test_gamma_transform_max_smaller_one(self):
        trafo = GammaCorrectionTransform(gamma=(0, 0.9))
        random.seed(0)
        rand_val = random.uniform(0, 0.9)
        random.seed(0)
        outp = trafo(**self.batch_dict)
        expected_out = self.batch_dict["data"].pow(rand_val)
        self.assertTrue((outp["data"] == expected_out).all())

    def test_gamma_transform_max_greater_one(self):
        # (1., 1.) allows switching cases but forces gamma to 1.
        trafo = GammaCorrectionTransform(gamma=(1., 1.))
        random.seed(0)
        for _ in range(5):
            outp = trafo(**self.batch_dict)
            self.assertTrue((outp["data"] == self.batch_dict["data"]).all())


if __name__ == '__main__':
    unittest.main()
