"""
correlation_spanning_tree.py
--------------

Graph reconstruction algorithm based on Mantegna, R. N. (1999). Hierarchical structure in 
financial markets. The European Physical Journal B-Condensed Matter and Complex Systems, 
11(1), 193-197. DOI https://doi.org/10.1007/s100510050929
https://link.springer.com/article/10.1007/s100510050929

author: Matteo Chinazzi
Submitted as part of the 2019 NetSI Collabathon.
"""

from .base import BaseReconstructor
import numpy as np
import networkx as nx
from scipy.sparse.csgraph import minimum_spanning_tree


class CorrelationSpanningTree(BaseReconstructor):
    def fit(self, TS, **kwargs):
        """Create a minimum spanning tree connecting the sensors.
        
        The empirical correlation matrix is used to first compute a distance matrix
        and then to create a minimum spanning tree connecting all the sensors in the data. 
        This method implements the methodology described in [1] and applied in the context 
        of creating a graph connecting the stocks of a portfolio of generated by looking
        at the correlations between the daily time series of stock prices.

        Params
        ------
        TS (np.ndarray): $N \\times L$ array consisting of $L$ observations 
                         from $N$ sensors.


        Returns
        -------
        G (nx.Graph): A reconstructed graph with $N$ nodes.

        Example
        --------

        import numpy as np
        import networkx as nx
        from matplotlib import pyplot as plt
        from netrd.reconstruction import CorrelationSpanningTree

        N = 25
        T = 300
        M = np.random.normal(size=(N,T))

        print('Create correlated time series')
        market_mode = 0.4*np.random.normal(size=(1,T))
        M += market_mode

        sector_modes = {d: 0.5*np.random.normal(size=(1,T)) for d in range(5)}
        for sector_mode, vals in sector_modes.items():
            M[sector_mode*5:(sector_mode+1)*5,:] += vals

        print('Link node colors to sectors')
        colors = ['b','r','g','y','m']
        node_colors = [color for color in colors for __ in range(5)]

            
        print('Network reconstruction step')
        cst_net = CorrelationSpanningTree()
        G = cst_net.fit(M)

        print('Plot reconstructed spanning tree')
        fig, ax = plt.subplots()
        nx.draw(G, ax=ax, node_color=node_colors)


        References
        ----------
        .. [1] Mantegna, R. N. (1999). Hierarchical structure in financial markets. 
        The European Physical Journal B-Condensed Matter and Complex Systems, 11(1), 193-197. 
        DOI https://doi.org/10.1007/s100510050929
        https://link.springer.com/article/10.1007/s100510050929


        """
        N, L = TS.shape

        C = np.corrcoef(TS) # Empirical correlation matrix
        
        D = np.sqrt(2*(1-C)) # Metric distance matrix

        MST = minimum_spanning_tree(D) # Minimum Spanning Tree

        G = nx.from_scipy_sparse_matrix(MST)

        self.results['graph'] = G

        return G

