# -*- coding: utf-8 -*-
"""
This module allows to visualize, plot and compute metrics on the infered gene regulatory networks, using built-in metrics in `networkx`_.

.. _scikit-learn:
    https://scikit-learn.org
"""
import pandas as pd
import numpy as np
import networkx as nx
import matplotlib.pyplot as plt
from scipy import stats
from grenadine.Inference.inference import rank_GRN
from sklearn.decomposition import PCA
from sklearn.metrics import roc_curve
from sklearn.metrics import roc_auc_score
from sklearn.metrics import auc
from sklearn.metrics import precision_recall_curve
from sklearn.metrics import precision_score
from sklearn.metrics import recall_score
from sklearn.metrics import accuracy_score
from sklearn.metrics import f1_score
from tqdm import tqdm
from seaborn import color_palette,clustermap,boxplot
from matplotlib.patches import Patch


def plot_scores(scores_, metrics_list = None):
    """
    Plots the evaluation scores distribution as boxplots

    Args:
        metrics_list (list): list of metrics names in scores_ to be represented
    """
    if metrics_list is None:
        metrics_list = [c for c in scores_.columns if "mean_test" in c]
    s = scores_[metrics_list]
    boxplot(s,
            orient='h',notch=True, showcaps=False,
            flierprops={"marker": "."},
            boxprops={"facecolor": (.1, .1, .1, .1)},
            medianprops={"color": "darkred", "linewidth": 3})
    plt.xlabel('Mean best model score',fontsize=10)
    plt.xlabel('')
    _=plt.yticks(fontsize=12)
    _=plt.xticks(fontsize=12)
    plt.xlim([-0.03,1.03])


def clustermap_with_classes(df,classes_row, classes_col,**clustermap_params):
    row_colors = None
    col_colors = None
    if classes_row is not None:
        unique_classes_row =  classes_row.unique()
        unique_classes_row.sort()
        colors_row = color_palette("Greens",n_colors=len(unique_classes_row))
        colors_map_row = {c:colors_row[i] for i,c in enumerate(unique_classes_row)} 
        row_colors = classes_row.map(colors_map_row)

    if classes_col is not None:
        unique_classes_col =  classes_col.unique()
        unique_classes_col.sort()
        colors_col = color_palette("Blues",n_colors=len(unique_classes_col))
        colors_map_col = {c:colors_col[i] for i,c in enumerate(unique_classes_col)} 
        col_colors = classes_col.map(colors_map_col)
    
    default_clustermap_params = {'center':0,
                                 'method':'ward',
                                 'xticklabels':False,
                                 'yticklabels':False,
                                 'figsize':(3,3),
                                }
    
    default_clustermap_params.update(clustermap_params)
    
    ax = clustermap(df,
                    row_colors=row_colors,
                    col_colors=col_colors,
                    **default_clustermap_params,
                    )
    if classes_col is not None:
        handles = [Patch(facecolor=colors_map_col[name]) for name in colors_map_col]
        legend_col =plt.legend(handles, 
                               colors_map_col, 
                               title="Column class",
                               bbox_to_anchor=(-0., 1),
                               bbox_transform=plt.gcf().transFigure,
                               loc='upper right')
        plt.gca().add_artist(legend_col)

    if classes_row is not None:
        handles = [Patch(facecolor=colors_map_row[name]) for name in colors_map_row]
        legend_row =plt.legend(handles, 
                               colors_map_row, 
                               title="Row class",
                               bbox_to_anchor=(-0., 0.6),
                               bbox_transform=plt.gcf().transFigure,
                               loc='upper right')
        plt.gca().add_artist(legend_row)
    return ax


def to_graph(ranking_df):
    """
    This method returns a directed graph generated by networkx
    """
    return(nx.from_pandas_edgelist(ranking_df,source='TF',target='TG',edge_attr='score'))

def draw_by(G, dico, maxi):
    colnodes = list()
    sizenodes = list()
    for node in G.nodes:
        colnodes.append(dico[node])
        sizenodes.append((120.*dico[node])/maxi)
    nx.draw(G, pos = nx.nx_agraph.graphviz_layout(G), node_color = colnodes, node_size = sizenodes, cmap = 'plasma')


def rich_club_coefficient(G,show=False):
    undir_G=nx.to_undirected(G)
    rich_club = nx.rich_club_coefficient(undir_G)
    if(show):
        plt.bar(*zip(*rich_club.items()))
        plt.xlabel('k minimum')
        plt.ylabel('Coefficient rich club')
        plt.show()
    return(rich_club)


def betweenness_centrality(G,show=False):
    between = nx.betweenness_centrality(G)
    if(show):
        plt.bar(*zip(*between.items()))
        plt.xlabel('gene')
        plt.ylabel('Betweeness centrality')
        plt.show()
    return(between)

def clustering_coefficient(G,show=False):
    undir_G=nx.to_undirected(G)
    clust = nx.clustering(undir_G)
    if(show):
        plt.bar(*zip(*clust.items()))
        plt.xlabel('gene')
        plt.ylabel('Clustering coefficient')
        plt.show()
    return(clust)

def closeness_centrality(G,show=False):
    undir_G=nx.to_undirected(G)
    closeness = nx.closeness_centrality(undir_G)
    if(show):
        plt.bar(*zip(*closeness.items()))
        plt.xlabel('gene')
        plt.ylabel('Closeness centrality')
        plt.show()
    return(closeness)
