# 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 the HCBM and the time series generated from HCBM methods
in data_generation/hcbm.py
"""

import unittest
import numpy as np
from mlfinlab.data_generation.hcbm import generate_hcmb_mat, time_series_from_dist


class TestHCBM(unittest.TestCase):
    """
    Tests the data generated from the HCBM module.
    """

    def setUp(self):
        """
        Sets the random number generator seed and parameters for functions.
        """

        np.random.seed(2814)

        self.t_samples = [4, 2, 1, 1]
        self.n_size = [10, 20, 17, 200]
        self.rho_lows = [0.1, 0.2, 0.9, 0.1]
        self.rho_highs = [0.9, 0.8, 0.95, 0.9]
        self.blocks = [3, 4, 10, 6]
        self.depths = [2, 2, 10, 3]
        self.permute = [False, True, False, False]

        self.ts_samples = [200, 100, 100, 500]
        self.dists = ["normal", "student", "student", "normal"]

        self.n_generated = len(self.t_samples)

    @staticmethod
    def _data_exists(mats):
        """
        Returns whether matrices exists and it is greater and 0
        """

        return mats is not None and len(mats) > 0

    @staticmethod
    def _data_is_symmetric(mats, dim):
        """
        Returns whether matrices generated are symmetrical.
        """

        result = True
        diag_rows, diag_cols = np.triu_indices(dim, k=1)

        for corr_mat in mats:
            result &= (corr_mat[diag_rows, diag_cols] == corr_mat[diag_cols, diag_rows]).all()

        return result

    @staticmethod
    def _data_correctly_bounded(mats, dim, rho_low, rho_high):
        """
        Returns whether matrix is correctly bounded.
        """

        diag_rows, diag_cols = np.diag_indices(dim)
        diag_is_ones = (mats[:, diag_rows, diag_cols] == 1).all()
        rows, cols = np.triu_indices(dim, 1)
        mat_is_lower_bounded = (mats[:, rows, cols] >= rho_low).all()
        mat_is_upper_bounded = (mats[:, rows, cols] <= rho_high).all()

        return diag_is_ones and mat_is_lower_bounded and mat_is_upper_bounded

    def test_data_returned_hcbm(self):
        """
        Tests that data generated from HCBM method exists for multiple dimensions and
        parameters.
        """

        for i in range(self.n_generated):
            hcbm_mat = generate_hcmb_mat(
                t_samples=self.t_samples[i],
                n_size=self.n_size[i],
                rho_low=self.rho_lows[i],
                rho_high=self.rho_highs[i],
                blocks=self.blocks[i],
                depth=self.depths[i],
                permute=self.permute[i],
            )
            self.assertTrue(self._data_exists(hcbm_mat))

    def test_correct_shapes_hcbm(self):
        """
        Tests that the data generated from the HCBM method exists for multiple dimensions and
        parameters.
        """

        for i in range(self.n_generated):
            hcbm_mat = generate_hcmb_mat(
                t_samples=self.t_samples[i],
                n_size=self.n_size[i],
                rho_low=self.rho_lows[i],
                rho_high=self.rho_highs[i],
                blocks=self.blocks[i],
                depth=self.depths[i],
                permute=self.permute[i],
            )
            self.assertTrue(hcbm_mat.shape == (self.t_samples[i], self.n_size[i], self.n_size[i]))

    def test_data_is_correctly_bounded_hcbm(self):
        """
        Tests that the data generated from the HCBM method is bounded to the range
        rho_low <= data <= rho_high and that its diagonal is always equal to 1s.
        """

        for i in range(self.n_generated):
            hcbm_mat = generate_hcmb_mat(
                t_samples=self.t_samples[i],
                n_size=self.n_size[i],
                rho_low=self.rho_lows[i],
                rho_high=self.rho_highs[i],
                blocks=self.blocks[i],
                depth=self.depths[i],
                permute=self.permute[i],
            )
            self.assertTrue(
                self._data_correctly_bounded(
                    hcbm_mat, self.n_size[i], self.rho_lows[i], self.rho_highs[i]
                )
            )

    def test_matrix_is_symmetrical_hcbm(self):
        """
        Tests that the data generated from the HCBM method are symmetrical.
        """

        for i in range(self.n_generated):
            hcbm_mat = generate_hcmb_mat(
                t_samples=self.t_samples[i],
                n_size=self.n_size[i],
                rho_low=self.rho_lows[i],
                rho_high=self.rho_highs[i],
                blocks=self.blocks[i],
                depth=self.depths[i],
                permute=self.permute[i],
            )
            self.assertTrue(self._data_is_symmetric(hcbm_mat, self.n_size[i]))

    @staticmethod
    def test_small_data_hcbm():
        """
        Tests that the HCBM method returns valid data.
        Valid data is defined as not all values equal to 0.
        """

        hcbm_sample = [[[1, 0.51696, 0.25, 0.25, 0.25, 0.25],
                        [0.51696, 1, 0.25, 0.25, 0.25, 0.25],
                        [0.25, 0.25, 1, 0.55194, 0.31992, 0.31992],
                        [0.25, 0.25, 0.55194, 1, 0.31992, 0.31992],
                        [0.25, 0.25, 0.31992, 0.31992, 1, 0.43508],
                        [0.25, 0.25, 0.31992, 0.31992, 0.43508, 1]]]
        hcbm_mat = generate_hcmb_mat(
            t_samples=1, n_size=6, rho_low=0.25, rho_high=0.75, blocks=2, depth=2, permute=False
        )
        np.testing.assert_array_almost_equal(hcbm_mat, hcbm_sample, 5)

    def test_data_returned_time_series(self):
        """
        Tests that data generated from HCBM method exists for multiple dimensions and
        parameters.
        """

        for i in range(self.n_generated):
            hcbm_mat = generate_hcmb_mat(
                t_samples=1,
                n_size=self.n_size[i],
                rho_low=self.rho_lows[i],
                rho_high=self.rho_highs[i],
                blocks=self.blocks[i],
                depth=self.depths[i],
                permute=self.permute[i],
            )

            time_series = time_series_from_dist(
                corr=hcbm_mat[0], t_samples=self.ts_samples[i], dist=self.dists[i]
            )
            self.assertTrue(self._data_exists(time_series))

    def test_correct_shapes_time_series(self):
        """
        Tests that the data generated from the HCBM method exists for multiple dimensions and
        parameters.
        """

        for i in range(self.n_generated):
            hcbm_mat = generate_hcmb_mat(
                t_samples=1,
                n_size=self.n_size[i],
                rho_low=self.rho_lows[i],
                rho_high=self.rho_highs[i],
                blocks=self.blocks[i],
                depth=self.depths[i],
                permute=self.permute[i],
            )

            time_series = time_series_from_dist(
                corr=hcbm_mat[0], t_samples=self.ts_samples[i], dist=self.dists[i]
            )
            self.assertTrue(time_series.shape == (self.ts_samples[i], self.n_size[i]))

    def test_small_data_time_series(self):
        """
        Tests that the time series method returns valid data.
        Valid data is defined as not all values equal to 0.
        """
        np.set_printoptions(precision=5)

        hcbm_mat = generate_hcmb_mat(
            t_samples=1, n_size=6, rho_low=0.25, rho_high=0.75, blocks=2, depth=2, permute=False
        )
        time_series = time_series_from_dist(corr=hcbm_mat[0], t_samples=6, dist="normal").round(5)
        flat_series = [item for sublist in time_series.values.tolist() for item in sublist]

        time_series_sample = [[0.98732, -0.22608, 0.67257, 1.01905, 0.49815, 0.27573],
                              [-0.49786, -0.25026, 0.96379, 0.1083, 0.90386, 2.23842],
                              [-0.65337, -1.29134, -1.17062, 0.10184, 0.418, -0.07461],
                              [0.00385, -1.2113, 1.10762, -0.20205, 0.9475, -0.23154],
                              [0.34948, 1.7678, 0.66294, 0.44735, 1.02399, 0.68903],
                              [-0.70259, -0.05277, -1.17347, -1.01258, -0.37828, -0.6414]]
        flat_sample = [item for sublist in time_series_sample for item in sublist]

        self.assertCountEqual(flat_series, flat_sample)

    def test_time_series_fails_on_wrong_dist(self):
        """
        Tests that time series generation fails when given a not supported distribution name.
        """
        self.assertRaises(ValueError, time_series_from_dist, corr=[[]], t_samples=10, dist="any_dist")
