# Copyright 2019, Hudson and Thames Quantitative Research
# All rights reserved
# Read more: https://github.com/hudson-and-thames/mlfinlab/blob/master/LICENSE.txt
"""
Tests the data generated by CorrGAN in data_generation/corrgan.py.
"""
import unittest
import os
import numpy as np
import tensorflow as tf
from mlfinlab.data_generation.corrgan import sample_from_corrgan


class TestCorrGAN(unittest.TestCase):
    """
    Test the data sampled from CorrGAN pre-trained modules.
    """

    def setUp(self):
        """
        Sets the file path for the corrgan generator model.
        """

        project_path = os.path.dirname(__file__)
        path = project_path + "/test_data"
        self.generator_path = path

        tf.random.set_seed(2814)

    def test_sample_returned(self):
        """
        Tests that data generated from CorrGAN exists for dimensions <= 10.
        """

        corrgan_samples = sample_from_corrgan(model_loc=self.generator_path, dim=10, n_samples=1)
        corrgan_samples_small = sample_from_corrgan(
            model_loc=self.generator_path, dim=5, n_samples=1
        )

        self.assertTrue(corrgan_samples is not None)
        self.assertTrue(len(corrgan_samples) > 0)

        self.assertTrue(corrgan_samples_small is not None)
        self.assertTrue(len(corrgan_samples_small) > 0)

    def test_correct_shapes(self):
        """
        Tests that the data generated from CorrGAN has the expected data shape.
        """

        corrgan_sample_10_1 = sample_from_corrgan(
            model_loc=self.generator_path, dim=10, n_samples=1
        )
        corrgan_sample_10_4 = sample_from_corrgan(
            model_loc=self.generator_path, dim=10, n_samples=4
        )
        corrgan_sample_2_2 = sample_from_corrgan(model_loc=self.generator_path, dim=2, n_samples=2)
        corrgan_sample_7_3 = sample_from_corrgan(model_loc=self.generator_path, dim=7, n_samples=3)

        self.assertTrue(corrgan_sample_10_1.shape == (1, 10, 10))
        self.assertTrue(corrgan_sample_10_4.shape == (4, 10, 10))
        self.assertTrue(corrgan_sample_2_2.shape == (2, 2, 2))
        self.assertTrue(corrgan_sample_7_3.shape == (3, 7, 7))

    def test_data_generated_is_correctly_bounded(self):
        """
        Tests that the data generated from CorrGAN is bounded to the range
        -1 <= data <= 1 which corresponds to any correlation matrix bounds,
        and that the diagonal is always equal to 1s.
        """

        dimension = 10
        corrgan_sample = sample_from_corrgan(
            model_loc=self.generator_path, dim=dimension, n_samples=10
        )
        diag_rows, diag_cols = np.diag_indices(dimension)

        self.assertTrue((corrgan_sample[:, diag_rows, diag_cols] == 1).all())
        self.assertTrue((corrgan_sample >= -1).all())
        self.assertTrue((corrgan_sample <= 1).all())

    def test_invalid_dimensions(self):
        """
        Tests that sampling CorrGAN does not support dimensions greater than
        the allowed range.
        """

        self.assertRaises(
            ValueError, sample_from_corrgan, model_loc=self.generator_path, dim=201, n_samples=10
        )
        self.assertRaises(
            ValueError, sample_from_corrgan, model_loc=self.generator_path, dim=1, n_samples=1
        )

    def test_small_data(self):
        """
        Tests that CorrGAN returns valid data.
        Valid data is defined as not all values equal to 0 (except for diagonal).
        """

        dimension = 4
        corrgan_sample = sample_from_corrgan(
            model_loc=self.generator_path, dim=dimension, n_samples=1
        )
        expected_sample = [[[1, 0.381, -0.1595, 0.15073],
                            [0.381, 1, 0.48509, 0.42076],
                            [-0.1595, 0.48509, 1, 0.64655],
                            [0.15073, 0.42076, 0.64655, 1]]]

        # Test values are valid.
        np.testing.assert_array_almost_equal(corrgan_sample, expected_sample, decimal=5)

    def test_matrix_is_symmetrical(self):
        """
        Tests that the matrix generated from CorrGAN is symmetrical.
        """

        dimension = 10
        corrgan_sample = sample_from_corrgan(
            model_loc=self.generator_path, dim=dimension, n_samples=4
        )
        diag_rows, diag_cols = np.diag_indices(dimension)
        for corr_mat in corrgan_sample:
            self.assertTrue(
                (corr_mat[diag_rows, diag_cols] == corr_mat[diag_cols, diag_rows]).all()
            )
