import os, sys
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import igraph as ig
from igraph import Graph
from .scii_circos import scii_interaction_matrix_process, scii_tensor_interaction_matrix_process, _cell_pairs_generate

def lr_link_graph_generate(interaction_matrix:pd.DataFrame,
                        cells:list=None,
                        separator:str="-",
                        cell_lr_separator:str="|",
                        reducer:int=3,
                        cell_colors="Set3",
                        lr_color_palette:str="tab20",
                        scii_tensor:bool=False,
                        ) -> Graph:
    '''
    Input:
        interaction_matrix: DataFrame. 
            Matrix contains interesting interaction in the TME region of interest, index represents cell-cell pairs
            while column represents ligand-receptor pairs
        cells: List.
            cell type names, which will be used to separate cells of the cell pair.
        separator: str.
            Separator used to combine ligand with receptor genes into LR pairs.
        cell_lr_separator: str.
            Separator used to combine cells with LR genes.
        reducer: int.
            The size of a vertex = weight of the vertex/reducer.
        cell_colors: str|dict.
            Define colors of difference cell type. User can set matplotlib's colormap (e.g. viridis, jet, tab10) or label_name -> color dict (e.g. dict(A="red", B="blue", C="green", ...)).
            lr_palette
        lr_color_palette: str|list.
            Define cmap of links between ligand with receptor.
        scii_tensor: bool.
            If the interaction_matrix was generated by scii, set the scii_tensor=True.
    Return:
        g: Graph containing all vertices(ligand or receptor) and their links information.
        
    '''

    #transfer matrix into igraph
    if scii_tensor:
        sectors, links, genes = scii_tensor_interaction_matrix_process(interaction_matrix, cells, separator=separator, cell_lr_separator=cell_lr_separator)
    else:
        sectors, links, genes, cells = scii_interaction_matrix_process(interaction_matrix, cell_lr_separator=cell_lr_separator)

    #generate color palette for cells and links
    if isinstance(cell_colors, str):
        cell_colors = dict(zip(cells, [tuple(x[0:3]) for x in plt.get_cmap(cell_colors, len(cells)).colors]))
    elif isinstance(cell_colors, dict):
        cell_colors = cell_colors
    else:
        raise Exception("cell_colors must be string representing the matplotlib's colormap or dict")

    vertices = list(sectors.keys())
    vertices_color = [cell_colors[x.split(cell_lr_separator)[0]] for x in vertices]
    vertices_dict = dict(zip(vertices, range(len(vertices))))
    edges_index = [(vertices_dict[x[0]], vertices_dict[x[1]]) for x in links]
    lr_links = set([(x[0].split(cell_lr_separator)[1], x[1].split(cell_lr_separator)[1]) for x in links])
    if isinstance(lr_color_palette, str):
        link_color_palette = [tuple(x[0:3]) for x in plt.get_cmap(lr_color_palette, len(lr_links)).colors]
    elif isinstance(lr_color_palette, list):
        if len(link_color_palette) < len(lr_links):
            raise Exception("the length of link_color_palette less than then number of lr_links: {0} < {1}".format(len(link_color_palette), len(lr_links)))
        link_color_palette = link_color_palette[0:len(lr_links)]
    link_colors = dict(zip(lr_links, link_color_palette))
    edges_color = [link_colors[(x[0].split(cell_lr_separator)[1], x[1].split(cell_lr_separator)[1])] for x in links]

    g = Graph(n=len(vertices), edges=edges_index, directed=True)
    g.vs['name'] = vertices
    g.vs['label'] = [x.split(cell_lr_separator, 1)[1] for x in vertices]
    g.vs['color'] = vertices_color
    g.vs['weight'] = [x/reducer if x>reducer*5 else 5 for x in sectors.values()]
    g.es['weight'] = [x[2] for x in links]
    g.es['color'] = edges_color
    return g

def cell_graph_generate(interaction_matrix:pd.DataFrame,
                        cells:list=None,
                        reducer:int=10,
                        cell_colors="Set3",
                        scii_tensor:bool=False,
                        ) -> Graph:
    '''
    Input:
        interaction_matrix: DataFrame. 
            Matrix contains interesting interaction in the TME region of interest, index represents cell-cell pairs
            while column represents ligand-receptor pairs
        cells: List.
            cell type names, which will be used to separate cells of the cell pair.
        separator: str.
            Separator used to combine ligand with receptor genes into LR pairs.
        cell_lr_separator: str.
            Separator used to combine cells with LR genes.
        reducer: int.
            The size of a vertex = weight of the vertex/reducer.
        cell_colors: str|dict.
            Define colors of difference cell type. User can set matplotlib's colormap (e.g. viridis, jet, tab10) or label_name -> color dict (e.g. dict(A="red", B="blue", C="green", ...)).
            lr_palette
        scii_tensor: bool.
            If the interaction_matrix was generated by scii, set the scii_tensor=True.
    Return:
        g: Graph containing all vertices(ligand or receptor) and their links information.
        
    '''

    #Normalize the value into 0~100
    scii_min, scii_max = interaction_matrix.min().min(), interaction_matrix.max().max()
    norm_interaction_matrix = ((interaction_matrix - scii_min)/(scii_max - scii_min)*100).apply(round).astype(int)

    if scii_tensor:
        if cells==None:
            raise Exception("When scii_tensor is True, the cells parameter must be given")
        cci_df = norm_interaction_matrix.sum(axis=1).to_frame()
        cell_pairs = _cell_pairs_generate(cells)
        cci_df[['sender', 'receiver']] = [cell_pairs[x] for x in cci_df.index]
        
    else:
        cci_df = norm_interaction_matrix.T.sum(axis=1).to_frame()
        cci_df[['sender', 'receiver']] = [list(x) for x in cci_df.index]    
    
    cells = list(set(cci_df['sender'].unique()) | set(cci_df['receiver'].unique()))
    cells_dict = dict(zip(cells, range(len(cells))))
    cci_df['sender_index'] = cci_df['sender'].map(cells_dict)
    cci_df['receiver_index'] = cci_df['receiver'].map(cells_dict)
    cci_df = cci_df.rename(columns = {'sender': 'from', 'receiver': 'to', 0: 'value'}).reset_index(drop=True)

    #generate color palette for cells and links
    if isinstance(cell_colors, str):
        cell_colors = dict(zip(cells, [tuple(x[0:3]) for x in plt.get_cmap(cell_colors, len(cells)).colors]))
    elif isinstance(cell_colors, dict):
        cell_colors = cell_colors
    else:
        raise Exception("cell_colors must be string representing the matplotlib's colormap or dict")

    vertices = cells
    vertices_color = [cell_colors[x] for x in cells]
    vertices_weight = cci_df['value'].groupby(cci_df['from']).sum()
    edges_index = [tuple(x) for x in cci_df[['sender_index', 'receiver_index']].values]
    edges_color = [cell_colors[cell] for cell in cci_df['from'].values]
    g = Graph(n=len(vertices), edges = edges_index, directed=True)
    g.vs['name'] = vertices
    g.vs['label'] = vertices
    g.vs['color'] = vertices_color
    g.vs['weight'] = [vertices_weight[x]/reducer if vertices_weight[x]>reducer*5 else 5 for x in vertices]
    g.es['weight'] = cci_df['value'].values
    g.es['color'] = edges_color

    return g

def cell_lr_grap_plot(g:Graph,
              separator:str='-',
              cell_lr_separator:str="|",
              layout_type:str='kk',
              save:str=None,
              vertex_label_angle=90,
              vertex_label_dist = 0,
              vertex_label_size = 8,
              edge_width=[0.5, 4],
              edge_curved=0.2,
              edge_arrow_size=10,
              edge_arrow_width=5,
              figsize:int=15,
              dpi:float=300,
              g_kwargs:dict={},
              ):
    '''
    Input:
        g: Graph.
            Graph object contians information of cell cell communication.
        separator: str.
            Separator used to combine ligand with receptor genes into LR pairs.
        cell_lr_separator: str.
            Separator used to combine cells with LR genes.
        layout_type: str.
            Layout style used to draw the graph. Same to the layout of igraph
        save: str.
            File name of the figure that will be saved.
        g_kwargs: dict.
            Dictionary containing parameters of igraph plot setting.
    Return:
        None
    '''
    
    cell_colors = dict(zip([x.split(cell_lr_separator)[0] for x in g.vs['name']], g.vs['color']))
    link_colors = dict(zip([f"{x.source_vertex['label']}{separator}{x.target_vertex['label']}" for x in g.es], g.es['color']))

    fig, ax = plt.subplots()
    layout = g.layout(layout_type)
    layout.rotate(0)
    #draw graph
    ig.plot(g,
            layout = layout,
            vertex_size=g.vs['weight'],
            vertex_label_angle=vertex_label_angle,
            vertex_label_dist = vertex_label_dist,
            vertex_label_size = vertex_label_size,
            edge_width=edge_width,
            edge_curved=edge_curved,
            edge_arrow_size=edge_arrow_size,
            edge_arrow_width=edge_arrow_width,
            target=ax,
            **g_kwargs
            )
    #generate legend
    cell_legend_handles = []
    for cell, color in cell_colors.items():
        handle = ax.scatter(
            [], [],
            s=100,
            facecolor=color,
            label=cell,
        )
        cell_legend_handles.append(handle)
    l1 = ax.legend(
        handles=cell_legend_handles,
        title='Cell Type',
        bbox_to_anchor=(1.0, 1.0),
        bbox_transform=ax.transAxes,
    )
    lr_legend_handles = []
    for lr, color in link_colors.items():
        handle = ax.scatter(
            [], [],
            s=100,
            facecolor=color,
            label=lr,
            marker = 's',
        )
        lr_legend_handles.append(handle)
    ax.legend(
        handles=lr_legend_handles,
        title='Ligand-Receptor',
        bbox_to_anchor=(1.0, 0.6),
        bbox_transform=ax.transAxes,
    )
    fig.gca().add_artist(l1)
    fig.set_size_inches(figsize, figsize)
    if save != None:
        fig.savefig(save, dpi=dpi)


def cell_graph_plot(g:Graph,
              layout_type:str='kk',
              save:str=None,
              vertex_label_angle=90,
              vertex_label_dist = 0,
              vertex_label_size = 8,
              edge_width=[0.5, 4],
              edge_curved=0.2,
              edge_arrow_size=10,
              edge_arrow_width=5,
              figsize:float=15,
              dpi:float=300,
              g_kwargs:dict={},
              ):
    '''
    Input:
        g: Graph.
            Graph object contians information of cell cell communication.
            layout_type: Layout style used to draw the graph. Same to the layout of igraph
        layout_type: str.
            Layout style used to draw the graph. Same to the layout of igraph
        save: str.
            File name of the figure that will be saved.
        g_kwargs: dict.
            Dictionary containing parameters of igraph plot setting.
    '''

    fig, ax = plt.subplots()
    layout = g.layout(layout_type)
    layout.rotate(0)
    #draw graph
    ig.plot(g,
            layout = layout,
            vertex_size=g.vs['weight'],
            vertex_label_angle=vertex_label_angle,
            vertex_label_dist = vertex_label_dist,
            vertex_label_size = vertex_label_size,
            edge_width=edge_width,
            edge_curved=edge_curved,
            edge_arrow_size=edge_arrow_size,
            edge_arrow_width=edge_arrow_width,
            target=ax,
            **g_kwargs
            )
    fig.set_size_inches(figsize, figsize)
    if save != None:
        fig.savefig(save, dpi=dpi)