import pandas as pd
from MergeTree import * # This is a custom package.

def DownloadDataset(dataset_url):
    ''' This function downloads the dataset from the given url and returns a pandas dataframe.
        Input: dataset_url
        Output: pandas dataframe
    '''
    df = pd.read_csv(dataset_url)
    return df

def GetLeafData(leaf_id):
    ''' This function returns a pandas dataframe containing all the data of a particular leaf. It first downloads the dataset from passiflora_aligned.csv and then returns the data of the leaf with the given leaf_id.
        Input: leaf_id
        Output: pandas series
    '''
    data = DownloadDataset('passiflora_aligned.csv')

    # Get the data of a single leaf with the given leaf_id
    leaf = data.iloc[leaf_id]
    return leaf

def GetCoordinates(data):
    ''' This function extracts the coordinates of the leaf from the given data.
        Input: data
        Output: dictionary of oute rcoordinates
    '''

    # Extracting the coordinates from the data in form of a list of tuples
    n_pos_list = list((data.iloc[i], data.iloc[i+1])
                      for i in range(6, (len(data)-1), 2))

    # Remove inner coordinates from the list
    for i in [1, 2, 3, 4]:
        del n_pos_list[1]

    # Converting the list of tuples into a dictionary with keys as the index of the tuple
    n_pos = dict(enumerate(n_pos_list))
    return n_pos

def LeaftoGraph(coordinate_dict):
    ''' This function visualizes the data in the form of a graph.
        Input: coordinate_dict
        Output: graph
    '''
    G = nx.Graph()

    # Adding nodes to the graph
    G.add_nodes_from(coordinate_dict.keys())

    for n, p in coordinate_dict.items():
        G.nodes[n]['pos'] = p

    # Adding edges to the graph
    nodes = list(G.nodes)
    edges = []
    for p_node, c_node in zip(nodes, nodes[1:]):
        edges.append((p_node, c_node))

    edges.append((nodes[0], nodes[-1]))

    G.add_edges_from(edges)

    return G

def DataVisualization(leaf_id):
    ''' This function visualizes the data in the form of a graph.
        Input: leaf_id
        Output: graph
    '''
    G = LeaftoGraph(GetCoordinates(GetLeafData(leaf_id)))
    # Drawing the graph
    nx.draw(G, nx.get_node_attributes(G, 'pos'), with_labels=True)

def ComputeMergeTree(leaf_id: int) -> nx.Graph:
    ''' This function computes the merge tree of a particular leaf in horizontal and vertical direction.
        Input: leaf_id
        Output: dictionary of height matrics in horizontal and vertical direction
    '''
    # Get the filtration in each direction
    hAngle = findFiltration(theta=0)
    vAngle = findFiltration(theta=90)

    G = LeaftoGraph(GetCoordinates(GetLeafData(leaf_id)))

    # Compute the merge tree in each direction
    hMergetree = buildMergeTreeAllPoints(
        g_orig=G, filtration=hAngle, show=False)
    vMergetree = buildMergeTreeAllPoints(
        g_orig=G, filtration=vAngle, show=False)

    # Compute the height matrix in each direction
    hHeight = getHeightMatrix(hMergetree)[0]
    vHeight = getHeightMatrix(vMergetree)[0]

    MTheight = [hHeight, vHeight]

    MT = dict(enumerate(MTheight))
    return MT