r"""
The downanlysis of cellphonedb
"""

import matplotlib.pyplot as plt
import pandas as pd
import numpy as np
import networkx as nx
import matplotlib.patches as mpatches
import scanpy as sc
import matplotlib
import anndata

kpy_install=False

def global_imports(modulename,shortname = None, asfunction = False):
    if shortname is None: 
        shortname = modulename
    if asfunction is False:
        globals()[shortname] = __import__(modulename)
    else:        
        globals()[shortname] = __import__(modulename)

def check_kpy():
    """
    
    """
    global kpy_install
    try:
        import ktplotspy as kpy
        kpy_install=True
        print('ktplotspy have been install version:',kpy.__version__)
    except ImportError:
        raise ImportError(
            'Please install the ktplotspy: `pip install ktplotspy`.'
        )

def cpdb_network_cal(adata:anndata.AnnData,pvals:list,celltype_key:str)->dict:
    r"""
    Calculate a CPDB (Cell Phone Database) network using gene expression data and return a dictionary of results.

    Arguments:
        adata: An AnnData object containing gene expression data.
        pvals: A list or array of p-values for each gene in the expression data.
        celltype_key: The name of the cell type key in adata.obs.

    Returns:
        cpdb_dict: A dictionary of results generated by `kpy.plot_cpdb_heatmap()` from the `ktplotspy` library.

    """
    check_kpy()
    global kpy_install
    if kpy_install==True:
        global_imports("ktplotspy","kpy")

    cpdb_dict=kpy.plot_cpdb_heatmap(
        adata = adata,
        pvals = pvals,
        celltype_key = celltype_key,
        figsize = (1,1),
        title = "",
        symmetrical = False,
        return_tables=True,
    )
    return cpdb_dict


def cpdb_plot_network(adata:anndata.AnnData,interaction_edges:pd.DataFrame,
                      celltype_key:str,nodecolor_dict=None,
                      edgeswidth_scale:int=10,nodesize_scale:int=1,
                      pos_scale:int=1,pos_size:int=10,figsize:tuple=(5,5),title:str='',
                      legend_ncol:int=3,legend_bbox:tuple=(1,0.2),legend_fontsize:int=10,
                     return_graph:bool=False)->nx.Graph:
    r"""
    Plot a network of interactions between cell types using gene expression data.

    Arguments:
        adata: An AnnData object containing gene expression data.
        interaction_edges: A DataFrame containing the edges of the network, including the source, target, and count of interactions.
        celltype_key: The name of the cell type key in adata.obs.
        nodecolor_dict: A dictionary mapping cell type names to node colors. If not provided, uses the colors specified in adata.uns.
        edgeswidth_scale: The scaling factor for edge width. Default is 10.
        nodesize_scale: The scaling factor for node size. Default is 1.
        pos_scale: The scaling factor for node positions. Default is 1.
        pos_size: The size of the node positions. Default is 10.
        figsize: The size of the plot. Default is (5,5).
        title: The title of the plot. Default is ''.
        legend_ncol: The number of columns in the legend. Default is 3.
        legend_bbox: The location of the legend. Default is (1,0.2).
        legend_fontsize: The font size of the legend. Default is 10.
        return_graph: If True, returns the network graph. Default is False.

    Returns:
        ax or G: The plot if `return_graph` is False, otherwise the network graph.

    """
    check_kpy()
    global kpy_install
    if kpy_install==True:
        global_imports("ktplotspy","kpy")

    #set Digraph of cellphonedb
    G=nx.DiGraph()
    for i in interaction_edges.index:
        G.add_edge(interaction_edges.loc[i,'SOURCE'],
                   interaction_edges.loc[i,'TARGET'],
                   weight=interaction_edges.loc[i,'COUNT'],)
    
    #set celltypekey's color
    if nodecolor_dict!=None:
        type_color_all=nodecolor_dict
    else:
        if '{}_colors'.format(celltype_key) in adata.uns:
            type_color_all=dict(zip(adata.obs[celltype_key].cat.categories,adata.uns['{}_colors'.format(celltype_key)]))
        else:
            if len(adata.obs[celltype_key].cat.categories)>28:
                type_color_all=dict(zip(adata.obs[celltype_key].cat.categories,sc.pl.palettes.default_102))
            else:
                type_color_all=dict(zip(adata.obs[celltype_key].cat.categories,sc.pl.palettes.zeileis_28))
    
    #set G_nodes_dict
    nodes=[]
    G_degree=dict(G.degree(G.nodes()))


    G_nodes_dict={}
    links = []
    for i in G.edges:
        if i[0] not in G_nodes_dict.keys():
            G_nodes_dict[i[0]]=0
        if i[1] not in G_nodes_dict.keys():
            G_nodes_dict[i[1]]=0
        links.append({"source": i[0], "target": i[1]})
        weight=G.get_edge_data(i[0],i[1])['weight']
        G_nodes_dict[i[0]]+=weight
        G_nodes_dict[i[1]]+=weight
        
    #plot
    fig, ax = plt.subplots(figsize=figsize) 
    pos = nx.spring_layout(G, scale=pos_scale, k=(pos_size)/np.sqrt(G.order()))
    p=dict(G.nodes)

    nodesize=np.array([G_nodes_dict[u] for u in G.nodes()])/nodesize_scale
    nodecolos=[type_color_all[u] for u in G.nodes()]
    nx.draw_networkx_nodes(G, pos, nodelist=p,node_size=nodesize,node_color=nodecolos)

    edgewidth = np.array([G.get_edge_data(u, v)['weight'] for u, v in G.edges()])/edgeswidth_scale
    nx.draw_networkx_edges(G, pos,width=edgewidth)


    #label_options = {"ec": "white", "fc": "white", "alpha": 0.6}
    #nx.draw_networkx_labels(G, pos, font_size=10,) #bbox=label_options)
    plt.grid(False)
    plt.axis("off")
    plt.xlim(-2,2)
    plt.ylim(-2,1.5)

    labels = adata.obs[celltype_key].cat.categories
    #用label和color列表生成mpatches.Patch对象，它将作为句柄来生成legend
    color = [type_color_all[u] for u in labels]
    patches = [mpatches.Patch(color=type_color_all[u], label=u) for u in labels ] 

    #plt.xlim(-0.05, 1.05)
    #plt.ylim(-0.05, 1.05)
    plt.axis("off")
    plt.title(title)
    plt.legend(handles=patches,
               bbox_to_anchor=legend_bbox, 
               ncol=legend_ncol,
               fontsize=legend_fontsize)
    if return_graph==True:
        return G
    else:
        return ax
    #return {'Graph':G,'ax':ax}

def cpdb_plot_interaction(adata:anndata.AnnData,cell_type1:str,cell_type2:str,
                          means:pd.DataFrame,pvals:pd.DataFrame,
                          celltype_key:str,genes=None,
                         keep_significant_only:bool=True,figsize:tuple = (4,8),title:str="",
                         max_size:int=1,highlight_size:float = 0.75,standard_scale:bool = True,
                         cmap_name:str='viridis',
                         ytickslabel_fontsize:int=8,xtickslabel_fontsize:int=8,title_fontsize:int=10)->matplotlib.axes._axes.Axes:
    r"""
    Plot a CellPhoneDB interaction.

    Arguments:
        adata: AnnData object containing the data.
        cell_type1: Name of cell type 1.
        cell_type2: Name of cell type 2.
        means: DataFrame containing the means for each interaction.
        pvals: DataFrame containing the p-values for each interaction.
        celltype_key: Key in adata.obs that contains the cell type information.
        genes: List of genes to include in the plot. If None, all genes are included.
        keep_significant_only: Whether to keep only significant interactions.
        figsize: Figure size.
        title: Title of the plot.
        max_size: Maximum size of the dots.
        highlight_size: Size of the dots for the highlighted interaction.
        standard_scale: Whether to standard scale the data.
        cmap_name: Name of the colormap to use.
        ytickslabel_fontsize: Fontsize of the yticks labels.
        xtickslabel_fontsize: Fontsize of the xticks labels.
        title_fontsize: Fontsize of the title.
    
    Returns:
        ax: Axes object containing the plot.

    """
    check_kpy()
    global kpy_install
    if kpy_install==True:
        global_imports("ktplotspy","kpy")

    fig=kpy.plot_cpdb(
        adata = adata,
        cell_type1 = cell_type1,
        cell_type2 = cell_type2, 
        means = means,
        pvals = pvals,
        celltype_key = celltype_key,
        genes = genes,
        keep_significant_only=keep_significant_only,
        figsize = figsize,
        title = "",
        max_size = max_size,
        highlight_size = highlight_size,
        standard_scale = standard_scale,
        cmap_name=cmap_name
    ).draw()
    
    #ytickslabels
    labels=fig.get_axes()[0].yaxis.get_ticklabels()
    plt.setp(labels, fontsize=ytickslabel_fontsize)

    #xtickslabels
    labels=fig.get_axes()[0].xaxis.get_ticklabels()
    plt.setp(labels, fontsize=xtickslabel_fontsize)

    fig.get_axes()[0].set_title(title,fontsize=title_fontsize)
    
    return fig.get_axes()[0]

def cpdb_submeans_exacted(means:pd.DataFrame,cell_names:str,cell_type:str='ligand')->pd.DataFrame:
    r"""
    Returns a subset of the means DataFrame for a given cell type and cell name.

    Arguments:
        means: DataFrame containing the means for each interaction.
        cell_names: Name of the cell type.
        cell_type: Whether the cell type is a ligand or a receptor.

    Returns:
        means: Subset of the means DataFrame for a given cell type and cell name.

    """
    if cell_type=='ligand':
        means_columns=means.columns[:11].tolist()+means.columns[means.columns.str.contains('{}\|'.format(cell_names))].tolist()
    elif cell_type=='receptor':
        means_columns=means.columns[:11].tolist()+means.columns[means.columns.str.contains('\|{}'.format(cell_names))].tolist()
    else:
        raise ValueError('cell_type must be ligand or receptor')
    return means.loc[:,means_columns]

def cpdb_interaction_filtered(adata:anndata.AnnData,cell_type1:str,cell_type2:str,
                              means:pd.DataFrame,pvals:pd.DataFrame,celltype_key:str,genes=None,
                         keep_significant_only:bool=True,figsize:tuple = (0,0),
                         max_size:int=1,highlight_size:float = 0.75,standard_scale:bool = True,cmap_name:str='viridis',)->list:
    r"""
    Returns a list of unique interaction groups from a CellPhoneDB analysis filtered by significance.

    Arguments:
        adata: Annotated data matrix containing normalized gene expression data.
        cell_type1: The cell type name for the first interacting partner.
        cell_type2: The cell type name for the second interacting partner.
        means: Dataframe containing the means of gene expression for each interaction.
        pvals: Dataframe containing the p-values for each interaction.
        celltype_key: The column name in the `adata.obs` dataframe containing cell type annotations.
        genes: List of gene names to filter on. If None, no filtering is applied.
        keep_significant_only: Whether to filter on significant interactions only (i.e. interactions with FDR < 0.05).
        figsize: Figure size. Width and height of the figure in inches. 
        max_size: Maximum size of the markers in the plot.
        highlight_size: Size of the markers for the highlighted interaction.
        standard_scale: Whether to standard scale the data.
        cmap_name: Name of the colormap to use.

    Returns:
        genes: List of unique interaction groups from a CellPhoneDB analysis filtered by significance.

    """
    check_kpy()
    global kpy_install
    if kpy_install==True:
        global_imports("ktplotspy","kpy")

    res=kpy.plot_cpdb(
        adata = adata,
        cell_type1 = cell_type1,
        cell_type2 = cell_type2, 
        means = means,
        pvals = pvals,
        celltype_key = celltype_key,
        genes = genes,
        keep_significant_only=keep_significant_only,
        figsize = figsize,
        title = "",
        max_size = max_size,
        highlight_size = highlight_size,
        standard_scale = standard_scale,
        cmap_name=cmap_name,
        return_table=True
    )

    return list(set(res['interaction_group']))