import numpy as np 

from .TopoMap import TopoMap

from .UnionFindComponents import UnionFind

class TopoTree(TopoMap):
    """Generate TopoTree representation of the input data points.

    Parameters
    ----------

    metric : {'cityblock', 'cosine', 'euclidean', 'l1', 'l2', 'manhattan',' braycurtis', 'canberra', \
                'chebyshev', 'correlation', 'dice', 'hamming', 'jaccard', 'kulsinski', 'mahalanobis', \
                'minkowski', 'rogerstanimoto', 'russellrao', 'seuclidean', 'sokalmichener', 'sokalsneath',\
                'sqeuclidean', 'yule'} or a callable, default='euclidean'
        Distance metric to compute the MST. Check out sklearn's pairwise_distances to see other options.

    approach: {'mlpack', 'nx', 'ANN'}, default='mlpack'
        Approach for computing the MST. If 'mlpack', mlpack's function emst is used. In this case, the only 
        metric available is euclidean. If 'nx', it uses sklearn's pairwise_distance function followed by 
        networkx's minimum_spanning_tree function. If 'ANN', it uses DiskANN's precomputed approximate 
        nearest neighbor instead of computing a MST. In this case, you need to provide the index_path.

    load_mst: Bool, default=False
        Option to load precomputed mst. If set to True, you need to provide the mst_path.

    mst_path: string, default=''
        Path to the precomputed mst. Only used if load_mst is True.

    index_path: string, default=''
        Path to index file generated by DiskANN. Only used if approach is set to 'ANN'.

    min_box_size: int, default=1
        Minumum number of points for a component to have an assigned box.
    """

    def __init__(self, 
                 metric='euclidean',
                 approach = 'mlpack', 
                 load_mst = False, 
                 mst_path = '',
                 index_path = '',
                 min_box_size=1,
                 ) -> None:
        self.metric = metric
        self.approach = approach
        self.load_mst  = load_mst
        self.mst_path = mst_path
        self.index_path = index_path
        self.min_box_size = min_box_size

        if metric!='euclidean' and approach=='mlpack':
            self.approach = 'nx'

        # Parameters to set if they were precomputed
        self.index = None
        self.mst = None
        self.sorted_edges = None
    
    def _merge_components_boxes(self, c_a, c_b, i_a, i_b, d, i):
        # If a has a box
        if not self.points_component[i_a] is None:

            # If b does not have a box -> Join b to a's box
            if self.points_component[i_b] is None:
                a_box_id = self.points_component[i_a]
                self.points_component[list(c_b)] = a_box_id
                self.components_info[a_box_id]['points'].extend(list(c_b))
                self.components_info[a_box_id]['size'] += len(c_b)
                self.components_info[a_box_id]['persistence'] = d
                self.components_info[a_box_id]['children'] += 1

            # If b has a box -> Create parent box
            else:
                parent_box_id = len(self.components_info)
                self.components_info.append({})
                self.components_info[parent_box_id]['id'] = parent_box_id
                self.components_info[parent_box_id]['size'] = len(c_a)+len(c_b)
                self.components_info[parent_box_id]['points'] = c_a
                self.components_info[parent_box_id]['points'].extend(c_b)
                self.components_info[parent_box_id]['persistence'] = d
                self.components_info[parent_box_id]['created_at'] = d
                self.components_info[parent_box_id]['children'] = 2

                a_box_id = self.points_component[i_a]
                b_box_id = self.points_component[i_b]
                self.components_info[a_box_id]['parent'] = parent_box_id
                self.components_info[a_box_id]['died_at'] = d
                self.components_info[a_box_id]['persistence'] = d-self.components_info[a_box_id]['created_at']
                self.components_info[b_box_id]['parent'] = parent_box_id
                self.components_info[b_box_id]['died_at'] = d
                self.components_info[b_box_id]['persistence'] = d-self.components_info[b_box_id]['created_at']

                persistence_density_a = self.components_info[a_box_id]['size']/d
                self.components_info[a_box_id]['persistence_density'] = persistence_density_a
                persistence_density_b = self.components_info[b_box_id]['size']/d
                self.components_info[b_box_id]['persistence_density'] = persistence_density_b

                self.points_component[list(c_a)] = parent_box_id
                self.points_component[list(c_b)] = parent_box_id
                
        # If a does not have a box
        else:
            # If b has a box -> Join a to b's box
            if not self.points_component[i_b] is None:
                b_box_id = self.points_component[i_b]
                self.points_component[list(c_a)] = b_box_id
                self.components_info[b_box_id]['points'].extend(list(c_a))
                self.components_info[b_box_id]['size'] += len(c_a)
                self.components_info[b_box_id]['persistence'] = d
                self.components_info[b_box_id]['children'] += 1

            # If none has box
            else:
                # If a and b merged is bigger than min -> Create box
                if len(c_a)+len(c_b) >= self.min_box_size:
                    new_box_id = len(self.components_info)
                    self.components_info.append({})
                    self.components_info[new_box_id]['id'] = new_box_id
                    self.components_info[new_box_id]['size'] = len(c_a)+len(c_b)
                    self.components_info[new_box_id]['points'] = c_a
                    self.components_info[new_box_id]['points'].extend(c_b)
                    self.components_info[new_box_id]['persistence'] = d
                    self.components_info[new_box_id]['created_at'] = d
                    self.components_info[new_box_id]['children'] = 0

                    self.points_component[list(c_a)] = new_box_id
                    self.points_component[list(c_b)] = new_box_id

    def _get_components(self):

        for i in range(len(self.sorted_edges)):
            # Get points from the edge
            i_a, i_b = self.sorted_edges[i][0], self.sorted_edges[i][1]

            # Distance between points
            d = self.sorted_edges[i][2]

            # Get components the points belong to
            root_a, root_b = self.uf.find(i_a), self.uf.find(i_b)
            c_a = self.uf.component[root_a].points_ids.copy()
            c_b = self.uf.component[root_b].points_ids.copy()

            # Merge components 
            self._merge_components_boxes(c_a, c_b, i_a, i_b, d, i)
            self.uf.union(i_a, i_b)

        return self.components_info
    
    def fit(self, X:np.ndarray):
        self.uf = UnionFind(len(X), compute_hulls=False)
        self.points_component = np.array([None for i in range(len(X))])
        self.components_info = []

        if self.mst is None:
            self.mst = self._compute_mst(X)

        if self.sorted_edges is None:
            self.sorted_edges = self._compute_ordered_edges(self.mst)

        return self._get_components()
    