import pandas as pd
import numpy as np 

from .utils import Transform, get_hull
from .UnionFindComponents import UnionFind

from .TopoMap import TopoMap
from .TopoTree import TopoTree

class HierarchicalTopoMap(TopoMap):
    """Generate HierarchicalTopoMap projection 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_points_component : int, default=2
        Minumum number of points for a component to have an assigned box.

    max_edge_length : float, default=-1 (ignore restriction)
        Optional parameter to control the maximum length for an edge to be processed.

    components_to_scale : list, default=[]
        List of component ids to be scaled. If it is empty, the method get_component_to_scale 
        will be used to select the components.

    max_scalar : float, default=20
        Maximum scalar alpha to scale a component.

    selection_method : {'k_components', 'min_persistence', 'min_size', 'max_persistence', \
                        'max_size'}, default='k_components'
        Method for selecting the components to be scale. Only used if components_to_scale is empty.

    k_components : int, default=5
        Number of components to select for scaling. The selection starts on the root of the tree
        and moves towards the leafs until we have selected k_components. If we reach a component
        wihtout childs, it is added to the list of selected components. Otherwise, we continue
        processing the component's childs until the end criteria is met.
        Only used if components_to_scale is empty and selection_method is 'k_components'.

    min_persistence_selection : float, default=0.1
        Minimum persistence a component must have to be selected. The selection starts on the root 
        of the tree and moves towards the leafs while the components have persistence above this value. 
        If we reach a component wihtout childs, it is added to the list of selected components. 
        If all childs of this component have persistence less then this value, the component
        is added to the list of selected components. Otherwise, we continue processing the component's 
        childs until there are no more components to explore.
        Only used if components_to_scale is empty and selection_method is 'min_persistence'.

    min_size_selection : float, default=0.1
        Minimum number of points a component must have to be selected. The selection starts on the root 
        of the tree and moves towards the leafs while the components have size above this value. 
        If we reach a component wihtout childs, it is added to the list of selected components. 
        If all childs of this component have size smaller then this value, the component is added to the 
        list of selected components. Otherwise, we continue processing the component's childs until 
        there are no more components to explore.
        Only used if components_to_scale is empty and selection_method is 'min_size'.

    max_persistence_selection : float, default=0.1
        Maximum persistence a component must have to be selected. The selection starts on the root 
        of the tree and moves towards the leafs until the components have persistent below this value. 
        If the component persistence below this value, it is added to the list of selected components. 
        Otherwise, we continue exploring the component's childs.
        Only used if components_to_scale is empty and selection_method is 'max_persistence'.

    max_size_selection : float, default=1
        Maximum number of points a component must have to be selected. The selection starts on the root 
        of the tree and moves towards the leafs until the components size is below this value. 
        If the component size below this value, it is added to the list of selected components. 
        Otherwise, we continue exploring the component's childs.
        Only used if components_to_scale is empty and selection_method is 'max_size'.

    verbose : bool, default=True
        Option to print information about the execution.
    """   

    def __init__(self,
                 metric='euclidean',
                 approach = 'mlpack', 
                 load_mst = False, 
                 mst_path = '',
                 index_path = '',
                 min_points_component=2,
                 max_edge_length=-1,
                 components_to_scale=[],
                 blowout_scalar=2,
                 max_scalar=20,
                 selection_method='k_components',
                 k_components=5,
                 min_persistence_selection=0.1,
                 min_size_selection=0,
                 max_persistence_selection=0.1,
                 max_size_selection=1,
                 verbose=True
                 ) -> None:

        # MST-related parameters
        self.verbose = verbose
        self.metric = metric
        self.approach = approach
        self.load_mst  = load_mst
        self.mst_path = mst_path
        self.index_path = index_path

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

        # Projection-related parameters
        self.min_points_component = min_points_component
        self.max_edge_length = max_edge_length
        self.components_to_scale = components_to_scale
        self.blowout_scalar = blowout_scalar
        self.max_scalar = max_scalar
        self.goal_length = max_edge_length

        # Options for get_component_to_scale
        self.selection_method = selection_method
        self.k_components = k_components
        self.min_persistence_selection = min_persistence_selection
        self.min_size_selection = min_size_selection
        self.max_persistence_selection = max_persistence_selection
        self.max_size_selection = max_size_selection

        # Parameters to set if they were precomputed
        self.index = None
        self.mst = None
        self.sorted_edges = None

    def _get_component_to_scale(self, 
                                X:np.ndarray,
                                df_comp=None, 
                                selection_method='k_components',
                                k_components=5,
                                min_persistence=0.1,
                                min_size=0,
                                max_persistence=0.1,
                                max_size=0):
        if df_comp is None:
            topotree = TopoTree(min_box_size=self.min_points_component)
            topotree.mst = self.mst
            topotree.sorted_edges = self.sorted_edges
            comp_info = topotree.fit(X)

            df_comp = pd.DataFrame.from_dict(comp_info)

        childs_map = {}
        for i in range(df_comp.shape[0]):
            childs_component = df_comp[df_comp['parent']==i].id.to_list()
            childs_map[i] = childs_component

        selected_comps = []

        if selection_method == 'k_components':
            ordered_comps = df_comp.sort_values(by=['size'], ascending=False).index.to_list()

            # Go down in the ordered list of components until we selected k
            i = 0
            while len(selected_comps) < k_components and i < len(ordered_comps):
                next_comp = ordered_comps[i]
                parent = df_comp.loc[next_comp,'parent']
                if parent in selected_comps:
                    selected_comps.remove(parent)
                selected_comps.append(next_comp)
                i += 1

            # Go up in the selected k and pick the parents that don't clash
            for i, c in enumerate(selected_comps):
                parents_clash = False
                parent_c = int(df_comp.loc[c,'parent'])
                while not parents_clash and not np.isnan(parent_c):
                    for j in range(len(selected_comps)):
                        if i==j:
                            continue
                        parent_j = df_comp.loc[selected_comps[j],'parent']
                        while not np.isnan(parent_j):
                            if parent_j == parent_c:
                                parents_clash = True
                                break
                            if not np.isnan(parent_j):
                                parent_j = df_comp.loc[parent_j,'parent']
                        if parents_clash:
                            break
                    if not parents_clash:
                        selected_comps = [parent_c if x==c else x for x in selected_comps]
                        c = parent_c
                        parent_c = int(df_comp.loc[c,'parent'])

        elif selection_method == 'min_persistence' or selection_method == 'min_size':
            if selection_method == 'min_persistence':
                feature = 'died_at'
                min_value = min_persistence
            else:
                feature = 'size'
                min_value = min_size

            parents = [df_comp.shape[0]-1]

            while len(parents) > 0:
                parents_aux = parents.copy()
                for parent in parents:
                    parents_aux.remove(parent)
                    if len(childs_map[parent]) > 0:
                        added_child = False
                        for child in childs_map[parent]:
                            if df_comp.loc[child,feature] > min_value:
                                parents_aux.append(child)
                                added_child = True
                        if not added_child:
                            selected_comps.append(parent)
                    else:
                        selected_comps.append(parent)
                parents = parents_aux

        elif selection_method == 'max_persistence' or selection_method == 'max_size':
            if selection_method == 'max_persistence':
                feature = 'died_at'
                max_value = max_persistence
            else:
                feature = 'size'
                max_value = max_size

            parents = [df_comp.shape[0]-1]

            while len(parents) > 0:
                parents_aux = parents.copy()
                for parent in parents:
                    parents_aux.remove(parent)
                    if df_comp.loc[parent,feature] < max_value:
                        selected_comps.append(parent)
                    else:
                        parents_aux.extend(childs_map[parent])
                parents = parents_aux
        
        else:
            raise ValueError("Selection method not recognized. Options: 'layers_down', 'n_components', 'min_persistence', 'min_size', 'max_persistence', 'max_size'")
        
        self.components_to_scale = selected_comps
        return selected_comps
    
    def _scale_component(self, component_id, i_point):
        avg_edge_comp = self.points_component_avg_edge[i_point]
        alpha = self.blowout_scalar*self.goal_length/avg_edge_comp
        alpha = min([alpha, self.max_scalar])

        comp_ids = self.components_info[component_id]['points']
        self.points_scaled[comp_ids] = True

        if self.verbose:
            print(f'Scalling component {component_id} - Scalar: {alpha:.3f}', end='')
            points = self.projections[comp_ids,:]
            area = (points[:,0].max() - points[:,0].min())*(points[:,1].max() - points[:,1].min())
            print(f' - initial area: {area:.3f}...', end='')

        t_scale = Transform(scalar=alpha)
        self.projections[comp_ids,:] = t_scale.scale(self.projections[comp_ids,:])

        if self.verbose:
            points = self.projections[comp_ids,:]
            area = (points[:,0].max() - points[:,0].min())*(points[:,1].max() - points[:,1].min())
            print(f' final area: {area:.3f}.')

        return alpha

    def _merge_components_boxes(self, c_a, c_b, i_a, i_b, d):
        # 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]['points'] = c_a.copy()
                self.components_info[parent_box_id]['points'].extend(c_b)
                self.components_info[parent_box_id]['size'] = len(c_a)+len(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]
                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]['average_edge'] = self.points_component_avg_edge[i_a]

                b_box_id = self.points_component[i_b]
                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]['average_edge'] = self.points_component_avg_edge[i_b]
                
                self.points_component[list(c_a)] = parent_box_id
                self.points_component[list(c_b)] = parent_box_id

                # Scale a and b
                if a_box_id in self.components_to_scale:
                    alpha_a = self._scale_component(a_box_id, i_a)
                    self.components_info[a_box_id]['alpha'] = alpha_a
                if b_box_id in self.components_to_scale:
                    alpha_b = self._scale_component(b_box_id, i_b)
                    self.components_info[b_box_id]['alpha'] = alpha_b

        # 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_points_component:
                    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]['points'] = c_a.copy()
                    self.components_info[new_box_id]['points'].extend(c_b)
                    self.components_info[new_box_id]['size'] = len(c_a)+len(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]

            if self.max_edge_length!=-1 and d > self.max_edge_length:
                if self.verbose:
                    print(f'[INFO] Max edge length hit. Distance: {d} | max_edge_length: {self.max_edge_length}')
                self.next_dist = d
                break

            # 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 boxes and scale if necessary
            self._merge_components_boxes(c_a, c_b, i_a, i_b, d)

            # Get points from edge and components
            p_a, p_b = self.projections[i_a,:], self.projections[i_b,:]
            proj_c_a, proj_c_b = self.projections[list(c_a)], self.projections[list(c_b)]

            simplices_a = self.uf.component[root_a].hull_simplices
            simplices_b = self.uf.component[root_b].hull_simplices

            # Rotate the first to be the topmost
            proj_c_a, edge_t = self._rotate_component(c_a, p_a, simplices_a, direction='top')
            # Rotate the second to be the bottomost
            proj_c_b, edge_b = self._rotate_component(c_b, p_b, simplices_b, direction='bottom')

            # Translate components
            proj_c_a = self._translate_component(proj_c_a, edge_t, to_point=[0,0])
            proj_c_b = self._translate_component(proj_c_b, edge_b, to_point=[0,d])

            self.projections[list(c_a), :] = proj_c_a
            self.projections[list(c_b), :] = proj_c_b

            avg_edge_a = self.points_component_avg_edge[i_a]
            avg_edge_b = self.points_component_avg_edge[i_b]
            avg_edge_merged = (len(c_a)*avg_edge_a+len(c_b)*avg_edge_b+d)/(len(c_a)+len(c_b)+1)
            self.points_component_avg_edge[list(c_a)] = avg_edge_merged
            self.points_component_avg_edge[list(c_b)] = avg_edge_merged

            # Merge components 
            self.uf.union(i_a, i_b, self.projections)

        if i == len(self.sorted_edges)-1:
            if self.verbose:
                print(f'[INFO] Number of edges hit. Edges processed: {i}')
            self.next_dist = d

        return self.projections

    def get_scaled_hulls(self):
        for c in self.components_to_scale:
            c_points_id = self.components_info[c]['points']
            c_hull = get_hull(self.projections[list(c_points_id),:])
            self.components_info[c]['hull'] = c_hull

        return self.components_info

    def fit_transform(self, X:np.ndarray):
        """Fit X into an embedded space and return that transformed output.

        Parameters
        ----------
        X : ndarray of shape (n_samples, n_features)

        Returns
        -------
        X_new : ndarray of shape (n_samples, n_components)
            Embedding of the training data in two-dimensional space.
        """
         
        n = len(X)

        self.projections = np.zeros(shape=(n, 2), dtype=np.float32)
        self.uf = UnionFind(n)
        self.points_component = np.array([None for i in range(n)])
        self.points_component_avg_edge = np.zeros(n)
        self.components_info = []
        self.points_scaled = np.zeros(shape=(n), dtype=bool)

        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)

        if self.goal_length==-1:
            biggest_edge = self.sorted_edges[-1][2]
            self.goal_length = biggest_edge

        if len(self.components_to_scale) == 0:
            self.components_to_scale = self._get_component_to_scale(X=X,
                                                                    selection_method=self.selection_method,
                                                                    k_components=self.k_components,
                                                                    min_persistence=self.min_persistence_selection,
                                                                    min_size=self.min_size_selection,
                                                                    max_persistence=self.max_persistence_selection,
                                                                    max_size=self.max_size_selection
                                                                   )

        self._get_components()

        self.get_scaled_hulls()

        return self.projections
    