# -*- coding: utf-8 -*-
"""
Created on Tue Mar 28 09:57:33 2017

@author: vguillem
@author: hmenager
@author: hjulienn
@author: clasry
"""

from pandas import DataFrame, read_hdf
# create (or open) an hdf5 file and opens in append mode
import numpy as np
import matplotlib

matplotlib.use("AGG")
import matplotlib.pyplot as plt
from matplotlib import colors
import matplotlib.patches as mpatches

def replaceZeroes(df):
    ids = np.where((df !=0) & np.isfinite(df))
    min_nonzero = np.min(df.values[ids])
    df.values[df.values == 0] = min_nonzero
    return df


def create_global_plot(work_file_path: str, global_plot_path: str):
    df = read_hdf(work_file_path, "SumStatTab")
    df[['JASS_PVAL', 'UNIVARIATE_MIN_PVAL']] = replaceZeroes(df[['JASS_PVAL', 'UNIVARIATE_MIN_PVAL']])

    df["-log10(Joint p-value)"] = -np.log10(df.JASS_PVAL)
    df["ind"] = range(len(df))
    df["absolute_position"] = range(len(df))
    # df['CHR'] = df['CHR'].apply(lambda x: int(x[3:]))
    df_grouped = df.groupby(("CHR"))
    m = 0

    fig = plt.figure(figsize=(30, 12))
    ax = fig.add_subplot(111)
    x_labels = []
    x_labels_pos = []

    colors = [
        "#1f77b4",
        "#ff7f0e",
        "#2ca02c",
        "#d62728",
        "#9467bd",
        "#8c564b",
        "#e377c2",
        "#7f7f7f",
        "#bcbd22",
        "#17becf",
        "#1f77b4",
        "#ff7f0e",
        "#2ca02c",
        "#d62728",
        "#9467bd",
        "#8c564b",
        "#e377c2",
        "#7f7f7f",
        "#bcbd22",
        "#17becf",
        "#1f77b4",
        "#ff7f0e",
    ]
    for num, (name, group) in enumerate(df_grouped):
        group = DataFrame(group)
        group["absolute_position"] = group["position"] + m
        m = np.max(group["absolute_position"])
        group.plot(
            kind="scatter",
            x="absolute_position",
            y="-log10(Joint p-value)",
            s=10,
            color=colors[num],
            ax=ax,
            lw=0
        )
        x_labels.append(name)
        x_labels_pos.append(
            (
                group["absolute_position"].iloc[-1]
                - (
                    group["absolute_position"].iloc[-1]
                    - group["absolute_position"].iloc[0]
                )
                / 2
            )
        )
    ax.set_xticks(x_labels_pos)
    ax.set_xticklabels(x_labels)
    ax.set_xlim([0, m])

    ax.set_xlabel("CHR", fontsize=16)
    ax.set_ylabel("-log10(p-value)", fontsize=16)
    ax.set_title("Joint test association results by region", fontsize=18)
    fig.savefig(global_plot_path, dpi=300)
    fig.clf()



def create_quadrant_plot(work_file_path: str,
 quadrant_plot_path: str, significance_treshold=5*10**-8):
    """
    Create a "quadrant" plot that represent the joint test pvalue versus the univariate test pvalue for the most significant SNPs by genomic region. The plot use a logarithmic scale.

    :param work_file_path: path to the worktable
    :type work_file_path: str
    :param quadrant_plot_path: path to the file where to store the plot results
    :type quadrant_plot_path: str
    """

    df = read_hdf(work_file_path, "Regions")
    df[['JASS_PVAL', 'UNIVARIATE_MIN_PVAL']] = replaceZeroes(df[['JASS_PVAL', 'UNIVARIATE_MIN_PVAL']])

    #df["-log10(Joint p-value)"] = -np.log10(df.JASS_PVAL)

    log10_pval = -np.log10(df[['JASS_PVAL', 'UNIVARIATE_MIN_PVAL']])

    pv_t = DataFrame(log10_pval, columns=['JASS_PVAL', 'UNIVARIATE_MIN_PVAL'])

    n = df.shape[0]
    transparency = 0.6
    pv_t["color"] = "grey"

    # blue: significant pvalues for  and univariate tests
    cond = df.signif_status=="Both"
    pv_t.loc[cond.values, "color"] = "#3ba3ec"
    b = cond.sum()
    # red: significant pvalues for  test only
    cond = df.signif_status=="Joint"
    pv_t.loc[cond.values, "color"] = "#f77189"
    r = cond.sum()
    # green: significant pvalues for univariate test only
    cond = df.signif_status=="Univariate"
    pv_t.loc[cond.values, "color"] = "#50b131"
    c = cond.sum()

    # grey: non significant pvalues
    cond = df.signif_status=="None"
    a = cond.sum()

    fig, ax = plt.subplots(figsize=(10, 5))

    plt.subplot(121)
    plt.scatter(pv_t.UNIVARIATE_MIN_PVAL, pv_t.JASS_PVAL, c=pv_t.color.tolist(), alpha=0.6, s=10)

    plt.axis([0, pv_t.UNIVARIATE_MIN_PVAL.max(), 0, pv_t.JASS_PVAL.max()])
    # axes abcisse et ordonnée à 8
    treshold = -np.log10(significance_treshold)
    plt.axvline(treshold, color="grey", linestyle="--")
    plt.axhline(treshold, color="grey", linestyle="--")

    # légendes abcisse et ordonnée
    plt.xlabel('-log10(P) for univariate tests', fontsize=12)
    plt.ylabel('-log10(P) for joint test', fontsize=12)

    # légendes différents groupes
    # green: significant pvalues for univariate test only
    green_patch = mpatches.Patch(
        color='#50b131', label='{} Significant pvalues for univariate test only'.format(c))
    # blue: significant pvalues for  and univariate tests
    blue_patch = mpatches.Patch(
        color='#3ba3ec', label='{} Significant pvalues for joint and univariate tests'.format(b))
    # red: significant pvalues for  test only
    red_patch = mpatches.Patch(
        color='#f77189', label='{} Significant pvalues for joint test only'.format(r))
    # grey: non significant pvalues
    grey_patch = mpatches.Patch(color='grey', label='{} Non significant pvalues'.format(a))

    lgd = plt.legend(handles=[red_patch, blue_patch, green_patch, grey_patch],
               bbox_to_anchor=(0.95, -0.25), loc='lower center', ncol=2, mode="expand", borderaxespad=0.)

    plt.subplot(122)
    plt.scatter(pv_t.UNIVARIATE_MIN_PVAL, pv_t.JASS_PVAL, c=pv_t.color.tolist(), alpha=0.6, s=10)
    # axes abcisse et ordonnee à 8
    plt.axvline(treshold, color="grey", linestyle="--")
    plt.axhline(treshold, color="grey", linestyle="--")
    # zoom on the square of region detected by jass:
    alim = np.ceil(pv_t.JASS_PVAL.loc[pv_t.color == "#f77189"].max() + 2)
    if np.isnan(alim):
        alim = 10
    plt.axis([0, alim, 0, alim])
    # légendes abcisse et ordonnee
    plt.xlabel('-log10(P) for univariate tests', fontsize=12)
    #plt.show()
    plt.savefig(quadrant_plot_path, dpi=600,  bbox_extra_artists=(lgd,), bbox_inches='tight')
    plt.clf()

    nb_omnibus = r
    nb_total = r + b + c

    return (nb_omnibus, nb_total)


def create_qq_plot(work_file_path: str, qq_plot_path: str):
    df = read_hdf(work_file_path, "SumStatTab")
    df[['JASS_PVAL', 'UNIVARIATE_MIN_PVAL']] = replaceZeroes(df[['JASS_PVAL', 'UNIVARIATE_MIN_PVAL']])

    pvalue = -np.log10(df.JASS_PVAL)
    # Cast values between 0 and 1, 0 and 1 excluded
    x = -np.log10(np.arange(1, pvalue.shape[0] + 1) / (pvalue.shape[0] + 2))
    y = np.sort(pvalue[:, 0])
    plt.scatter(x[::-1], y, s=5)
    lambda_value = np.median(y) / np.median(x)
    x_1 = np.linspace(0, 6)
    y_1 = lambda_value * x_1
    x_2 = np.linspace(0, 6)
    plt.plot(x_1, y_1, c="red")
    plt.plot(x_2, x_2)
    plt.savefig(qq_plot_path, dpi=600)
