#===============================================================================
# collect.py
#===============================================================================

"""Calculate and plot collection curves"""

# Imports ======================================================================

import os.path
import pandas as pd
import seaborn as sns
from collections import Counter
import numpy as np
from math import prod, floor
from itertools import cycle
from scipy.stats import hypergeom



# Constants ====================================================================

COL_COLOR_PALETTE = 'mako_r'




# Functions ====================================================================

def col_values_conf(pk_results):
    pkr = pk_results
    sc_matrix = np.zeros((pkr.number_of_genomes,pkr.number_of_genomes))
    for _, _, score in pkr:
        sc_matrix[sum(score) - 1,:] += np.array(score[::-1])
    sc_df = pd.DataFrame(sc_matrix, index=range(1, pkr.number_of_genomes+1),
        columns=tuple(str(os.path.basename(g)).replace('.fasta.gz', '').replace('.fa.gz', '')
                      for g in reversed(pkr.genomes)))
    g = len(sc_df.columns)
    coef_matrix_total = np.array([[prod((g-s-i)/(g-1-i) for i in range(n)) for n in range(g)]
                            for s in range(1,g+1)])
    coef_matrix_core = np.flip(coef_matrix_total.copy(), axis=0)
    for n in range(1,g):
        coef_matrix_total[:,n] += coef_matrix_total[:,n-1]
    col_df_core = pd.DataFrame(sc_df.transpose().dot(coef_matrix_core).transpose())
    col_df_total = pd.DataFrame(sc_df.transpose().dot(coef_matrix_total).transpose())
    col_df = pd.concat((col_df_total, col_df_core))
    col_df['n_genomes']= tuple(range(1,g+1))*2
    plotting_data = col_df.melt(
        id_vars=['n_genomes'], value_name='n_kmers')
    del plotting_data['variable']
    plotting_data['sequence'] = (('total',) * g + ('core',) * g) * g
    return plotting_data


def col_values(pk_results, contours=None):
    """Calculate collection curve

    Parameters
    ----------
    pk_results
        a PKResults object
    contours
        if not None, an iterable of integers between 0 and 100

    Yields
    -------
        collection curve values
    """

    pkr = pk_results
    g = pkr.number_of_genomes
    score_dist = Counter(sum(score) for _, _, score in pkr)
    for n_genomes in range(1, g+1):
        yield n_genomes, sum((1-prod((g-s-n)/(g-n) for n in range(n_genomes)))*score_dist[s]
            for s in range(1, g+1)), 'total'
        if contours:
            for c in contours:
                yield n_genomes, sum((hypergeom.sf(floor(c/100*n_genomes), g, s, n_genomes))*score_dist[s]
                    for s in range(1, g+1)), f'{c}%'
        yield n_genomes, sum(prod((s-n)/(g-n) for n in range(n_genomes))*score_dist[s]
            for s in range(1, g+1)), 'core'


def col_plot(plotting_data, output, title: str = 'Collection curve',
             linewidth: int = 3, palette=COL_COLOR_PALETTE, alpha=1, width=4,
             height=3, legend_loc='best'):
    ax = sns.lineplot(x='n_genomes', y='n_kmers', hue='sequence',
                      data=plotting_data, linewidth=linewidth,
                      palette=palette, alpha=alpha)
    ax.set_title(title)
    ax.set_ylim(bottom=0)
    if legend_loc == 'outside':
        leg = ax.legend(bbox_to_anchor=(1.02, 1), loc='upper left',
                        borderaxespad=0)
    else:
        leg = ax.legend(loc=legend_loc)
    for line in leg.get_lines():
        line.set_linewidth(linewidth)
        line.set_alpha(alpha)
    fig = ax.get_figure()
    fig.set_figwidth(width)
    fig.set_figheight(height)
    fig.tight_layout()
    fig.savefig(output)
    fig.clf()


def collect(pk_results, output=None, title: str = 'Collection curve', linewidth: int = 3,
               palette=COL_COLOR_PALETTE, alpha=1, width=4, height=3,
               conf=False, contours=False, legend_loc='best'):
    if conf:
        col_df = col_values_conf(pk_results)
    else:
        col_df = pd.DataFrame(col_values(pk_results, contours=contours),
                              columns=('n_genomes', 'n_kmers', 'sequence'))
    if output:
        col_plot(col_df, output, title=title, linewidth=linewidth,
                 palette=(palette if contours else sns.color_palette(palette, n_colors=2)), alpha=alpha,
                 width=width, height=height, legend_loc=legend_loc)
    return col_df
