import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.container as container
import pbwrap.plot.violin as violin
import functools
from itertools import cycle

from .utils import make_color_frame, get_markers_generator, get_colors_list

def distribution_super_plot(data, ax, title=None, xlabel=None, ylabel=None, showextrema=False) :
    """
    Distribution plot for pd.Series data.
    Series is expected to have a multi-index from 1 to 3 dimensions and Series values are expected to be lists.

    Super plot is returned as plt.Axes

    """
    if type(data) != pd.Series : raise TypeError('data argument passed is of type {0} should be "pd.Series".'.format(type(data)))
    if type(ax) != plt.Axes : raise TypeError('ax argument passed is of type {0} should be "plt.Axes".'.format(type(ax)))

    level = len(data.index.names)

    if level == 1 :
        ax, legend = _distribution_lvl1(data, ax, showextrema=showextrema)
    
    elif level == 2 :
        ax, legend,_ = _distribution_lvl2(data, ax, showextrema=showextrema)

    elif level == 3 :
        ax, legend = _distribution_lvl3(data, ax, showextrema=showextrema)
    
    else : raise LevelError("Unsupported number of dimension in index (should be between 1 and 3).")
    
    legend_size = len(legend[0])
    if legend_size > 5 :
        ncols= (legend_size // 5) + 1
        legend = ax.legend(*legend, ncols=ncols)
    else : legend = ax.legend(*legend)
    
    if type(title) != type(None) : ax.set_title(title)
    if type(xlabel) != type(None) : ax.set_xlabel(xlabel)
    if type(ylabel) != type(None) : ax.set_ylabel(ylabel)

    return ax


class LevelError(IndexError) :
    pass

def _distribution_lvl1(data: pd.Series, ax: plt.Axes, showextrema=False) :

    multi_index = data.index.names
    if len(multi_index) != 1 : raise LevelError("_distribution_lvl1 was called but multi-index dimension does not match.")
    if type(data.iat[0]) != list : raise ValueError("for distribution plot a list is expected for each distributions to plot. {0} was found.".format(type(data.iat[0])))

    colors = get_colors_list(len(data))
    labels = list(data.index)

    ax = violin.violin_plot(
        ax=ax,
        distributions=data,
        labels= labels,
        colors= colors,
        showmeans= True,
        mean_size= 90,
        showextrema=showextrema,
    )

    legend_prep =(
        [plt.scatter(0,0, c= 'white', edgecolors= 'black')],
        ['distribution means']
        )

    return ax, legend_prep

def _distribution_lvl2(data: pd.Series, ax : plt.Axes, alpha= 0.6, showextrema=False, show_distribution_size=True) :
    
    multi_index = data.index.names
    if len(multi_index) != 2 : raise LevelError("_distribution_lvl2 was called but multi-index dimension does not match.")
    data = data.sort_index()
    measure = data.name
    distributions: pd.Series = data.reset_index(drop=False).groupby(multi_index[:1])[measure].apply(list)
    labels_lvl2 = list(distributions.index)
    if show_distribution_size :
        for index, label in enumerate(labels_lvl2) : 
            labels_lvl2[index] = label + '\n' + str([len(dis) for dis in distributions.iat[index]])
    labels_lvl1 = list(data.index.get_level_values(1).unique())

    colors_df = make_color_frame(labels= labels_lvl1)

    colors = list(pd.merge(data.reset_index(), colors_df, left_on= multi_index[1], right_on='labels').sort_values(multi_index)['colors'])

    ax, positions = violin.violin_plot(
        ax=ax,
        distributions=distributions,
        labels= labels_lvl2,
        colors=colors,
        alpha=alpha,
        linewith=1,
        showmeans= True,
        mean_size= 90,
        showextrema=showextrema,
        multi_violin_plot= True,
        return_position=True,
    )

    legend_prep = (
        [plt.scatter(0,0, c= 'white', edgecolors= 'black')] + [plt.bar(0,0, color= color) for color in colors_df['colors']],
        ['distribution means'] + list(colors_df.index)        
    )


    return ax, legend_prep, positions


def _extract_color_from_legend(legend_prep, remove_first=False) :
    if remove_first :
        legend_prep = (legend_prep[0][1:], legend_prep[1][1:])
    
    artists: list[container.BarContainer] = legend_prep[0]
    colors = [(artist.patches[0].get_facecolor(),) for artist in artists]

    color_df = pd.DataFrame(data= colors, index= legend_prep[1], columns= ['colors'])
    return color_df

def _distribution_lvl3(data: pd.Series, ax: plt.Axes, showextrema) :
    """

    yaxis = measure_value

    level1 = x-axis (axis label)
    level2 = distribution colors (legend)
    level3 = smaller scatter points within distribution whith different shapes (legend)

    """
    multi_index = data.index.names
    if len(multi_index) != 3 : raise LevelError("_distribution_lvl3 was called but multi-index dimension does not match.")
    data = data.sort_index()
    measure = data.name

    #Levevel 1&2
    list_flattener = functools.partial(sum, start=[])
    data_distribution_lvl2: pd.Series = data.reset_index(drop=False).groupby(multi_index[:2])[measure].apply(list).apply(list_flattener)
    ax, legend_lvl2, positions = _distribution_lvl2(
        data_distribution_lvl2, 
        ax,
        alpha=0.4,
        showextrema=showextrema
        )

    #Level 3
    index_lvl3 = data.index.get_level_values(2).unique()

    gen = get_markers_generator()
    next(gen) #skipping 'o' that is used for means of distributions

    marker_list = [next(gen) for i in range(len(index_lvl3))]
    marker_generator = cycle(marker_list)

    #This handles if not all replicates are presents in all rna/treatments
    lvl0_idx = list(data.index.get_level_values(0))
    lvl1_idx = list(data.index.get_level_values(1))
    lvl2_idx_unique = list(data.index.get_level_values(2).unique())

    lvl2_idx = lvl2_idx_unique * len(lvl0_idx)
    lvl0_idx *=  len(lvl2_idx_unique)
    lvl1_idx *=  len(lvl2_idx_unique)

    extended_lvl3_index = pd.MultiIndex.from_tuples(list(
        zip(lvl0_idx, lvl1_idx, lvl2_idx)
    )).drop_duplicates()

    data = data.reindex(extended_lvl3_index).sort_index()

    for position, index_lvl2 in zip(positions, data_distribution_lvl2.index) :
        for data_lvl_3, marker in zip(data.loc[index_lvl2], marker_generator) :

            if isinstance(data_lvl_3, float) : 
                continue
            
            data_lvl_3 = [np.mean(data_lvl_3)]
            x = np.array([position] * len(data_lvl_3), dtype=float)

            ax.scatter(
                x= x,
                y= data_lvl_3,
                s = 25,
                marker = marker,
                c=(0,0,0,0),
                edgecolors='black',
                linewidths=1,
                alpha= 1,
            )

    legend_lvl3 = (
        [plt.scatter(0,0, c= 'black', edgecolors= 'black', marker= marker) for marker in marker_list],
        [str(index) for index in lvl2_idx_unique]
    )

    legend_prep = [lvl2 + lvl3 for lvl2,lvl3 in zip(legend_lvl2,legend_lvl3)]

    return ax, legend_prep

def violin_plot(
        ax: plt.Axes, distributions, labels=None, sub_labels=None, colors=None, xlabel=None, ylabel=None, title=None, y_axis= None,
        linewith = 2, line_color = 'black', alpha= 0.6,
        vertical_plot=True, showmedians= True, showextrema = False, showmeans=False, mean_size = 45, multi_violin_plot= False
                ) :

    """
    Basis for violin plots.
    """

    #Check parameters

    check_parameter(ax = plt.Axes, xlabel = (type(None), str), ylabel = (type(None), str), vertical_plot= bool, title = (type(None),str), linewith = (int, float), line_color= str, y_axis= (type(None),tuple,list), multi_violin_plot= bool)
    if multi_violin_plot :
        for distrib in distributions :
            check_parameter(distrib = (list, tuple, np.ndarray, pd.DataFrame, pd.Series))

    if type(labels) == type(None) : pass
    elif len(distributions) != len(labels) : raise ValueError("Length of labels and distributions must match.")

    if type(colors) == type(None): pass
    elif len(distributions) != len(colors) and not multi_violin_plot: raise ValueError("Length of colors and distributions must match.")
    
    if type(y_axis) == type(None) : pass
    elif len(y_axis) != 2 : raise ValueError("y_axis is expected to be of length 2 : (ymin, ymax).")
    
    if multi_violin_plot :
        if type(colors) != type(None) :
            if len(colors) == len(distributions) :
                new_colors = []
                for color,distrib in zip(colors, distributions) :
                    new_colors.extend([color]*len(distrib))
                colors = new_colors

        max_individual_violin_number = max([len(distrib) for distrib in distributions]) + 1 #Is the maximum number of violin plotted for one set.
        positions, ticks_positions = multi_violin_plot_positions(distributions)
        distributions = list(chain(*distributions))
        assert len(distributions) == len(positions), "AssertionError : multi_distributions wrongly flattened : positions : {0}, distributions {1}".format(len(positions), len(distributions))

        #colors
        if type(colors) != type(None) :
            if len(colors) == len(distributions) : pass
            else : raise ValueError("Length of colors must either match length of distributions or the number of element in distributions")
    
    else :
        positions = np.arange(1, len(distributions) + 1)
        ticks_positions = np.arange(1, len(distributions) + 1)
        max_individual_violin_number = 1
    #Plot
    violin_plot = ax.violinplot(
        distributions,
        positions=positions, 
        vert= vertical_plot, 
        showmedians=showmedians,
        showextrema= showextrema,
        )

    if type(labels) == type(None) :
        labels = np.arange(1, len(distributions) + 1)
    
    xticks = ax.set_xticks(ticks_positions, labels=labels)


    ax.set_xlim(0.25, len(labels) * max_individual_violin_number + 0.75)
    if type(colors) == type(None) : colors = get_colors_list(len(violin_plot['bodies']))
    for violin, color in zip(violin_plot['bodies'], colors) :
        violin.set_facecolor(color)
        violin.set_alpha(alpha)

    for collection_name in ['cbars', 'cmins', 'cmaxes', 'cmedians'] :
        collection = violin_plot.get(collection_name)
        if type(collection) != type(None) :    
            collection.set_color(line_color)
            collection.set_linewidth(linewith)
    
    if type(y_axis) != type(None) :
        axis = list(ax.axis())
        if y_axis[0] != None : axis[2] = y_axis[0]
        if y_axis[1] != None : axis[3] = y_axis[1]
        ax.axis(axis)
    else : axis = ax.axis()

    if type(sub_labels) != type(None) :
        if len(sub_labels) != len(distributions) : raise ValueError("Length of sub_labels must match number of violins to plot.")
        for sub_label, x_position in zip(sub_labels, positions) :
            ax.text(x = x_position, y= axis[2], s= sub_label, ha= 'center')
    
    if showmeans :
        means = [np.mean(distrib) for distrib in distributions]
        ax.scatter(positions, means, c= colors, s= mean_size, linewidths=0.5, edgecolors='black')

    if type(xlabel) != type(None) : ax.set_xlabel(xlabel)
    if type(ylabel) != type(None) : ax.set_ylabel(ylabel)
    if type(title) != type(None) : ax.set_title(title)

    return ax

def multi_violin_plot_positions(distributions) :

    max_individual_violin_number = max([len(distrib) for distrib in distributions]) + 1#Is the maximum number of violin plotted for one set.

    positions = []
    ticks_positions = []
    for distrib_number, distrib in enumerate(distributions) :
        positions.extend(list(
            np.arange(1, len(distrib) + 1) + (distrib_number * max_individual_violin_number) if len(distrib) > 1 
            else [distrib_number * max_individual_violin_number + (max_individual_violin_number-1)/2 + 1]
        ))

        ticks_positions.append(
            distrib_number * max_individual_violin_number + (len(distrib)-1)/2 + 1 if len(distrib) > 1
            else distrib_number * max_individual_violin_number + (max_individual_violin_number-1)/2 + 1
        )

    return positions, ticks_positions
