# Copyright 2019, Hudson and Thames Quantitative Research
# All rights reserved
# Read more: https://github.com/hudson-and-thames/mlfinlab/blob/master/LICENSE.txt
"""
Tests hierarchical clustering methods located in clustering/hierarchical_clustering.py
"""

import unittest
import os
import numpy as np
import pandas as pd
from mlfinlab.clustering import optimal_hierarchical_cluster


class TestHierarchicalClustering(unittest.TestCase):
    """
    Test that the hierarchical clustering works.
    """

    def setUp(self):
        """
        Sets the random number generator seed.
        """
        np.random.seed(2814)
        project_path = os.path.dirname(__file__)
        stock_prices = pd.read_csv("{}/test_data/stock_prices.csv".format(project_path), index_col=0)
        stock_returns = stock_prices.pct_change()[1:]
        self.corr_mat = stock_returns.corr()

    def test_data_returned(self):
        """
        Tests that the data generated by the optimal clustering method is returned.
        """
        # Test input as numpy array.
        cluster_mat = optimal_hierarchical_cluster(self.corr_mat.values)
        self.assertTrue(cluster_mat is not None and len(cluster_mat) > 0)

        # Test input as pandas dataframe.
        cluster_mat = optimal_hierarchical_cluster(self.corr_mat)
        self.assertTrue(cluster_mat is not None and len(cluster_mat) > 0)

    def test_data_valid(self):
        """
        Tests that the data generated by the optimal clustering method is valid.
        """

        dim = self.corr_mat.shape[0]
        cluster_mat = optimal_hierarchical_cluster(self.corr_mat)
        diag_rows, diag_cols = np.diag_indices(dim)

        diag_is_ones = (cluster_mat[diag_rows, diag_cols] == 1).all()
        rows, cols = np.triu_indices(dim, 1)
        mat_is_lower_bounded = (cluster_mat[rows, cols] >= np.min(self.corr_mat.values) - np.finfo(np.float).eps).all()
        mat_is_upper_bounded = (cluster_mat[rows, cols] <= (np.max(self.corr_mat.values) + np.finfo(np.float).eps)).all()

        self.assertTrue(diag_is_ones and mat_is_lower_bounded and mat_is_upper_bounded)
        self.assertRaises(AssertionError, np.testing.assert_array_almost_equal, self.corr_mat.values, cluster_mat, 5)
