# -*- coding: utf-8 -*-
"""
This contains all functions for accessing the "worktable" hdf5 file.
All functions either create or read a worktable at a specific path location.
@author: vguillem, hmenager, hjulienne
"""
import logging
import os
import importlib

from typing import List

# from dask.dataframe import read_hdf as dask_read_hdf
from pandas import HDFStore, DataFrame, concat, read_hdf, read_csv, Series, Index

# create (or open) an hdf5 file and opens in append mode
import numpy as np
import scipy.stats as spst

import tables
import warnings

warnings.filterwarnings("ignore", category=tables.NaturalNameWarning)

from ..config import config
from .stats import (
    make_stat_computer_nopattern,
    make_stat_computer_pattern,
    make_stat_computer_nan_dumb,
)

def choose_stat_function(smart_na_computation, optim_na, function_name, stat_function, sub_cov, **kwargs):
    if smart_na_computation:
        #If stat is sumz use normal computer even with na
        if function_name == "omnibus_stat":
            if optim_na:
                stat_compute = make_stat_computer_pattern(sub_cov, stat_function)
            else:
                stat_compute = make_stat_computer_nan_dumb(sub_cov, stat_function)
        else:
            if function_name == "meta_analysis":
                stat_compute = make_stat_computer_nopattern(sub_cov, stat_function, **kwargs)
            elif function_name=="sumz_stat":
                loading_file = kwargs.get('loadings', None)
                if loading_file is None:
                    #Default loadings would be one for every phenotypes
                    stat_compute = make_stat_computer_nopattern(sub_cov, stat_function)
                else:
                    loadings = read_csv(loading_file, index_col=0)
                    loadings = loadings.iloc[:,0]
                    stat_compute = make_stat_computer_nopattern(sub_cov, stat_function,loadings=loadings )
            else:
                stat_compute = make_stat_computer_nopattern(sub_cov, stat_function)
    else:
        stat_compute = make_stat_computer_nopattern(sub_cov, stat_function)


    return stat_compute

def add_signif_status_column(region_sub_tab, significance_treshold=5*10**-8):


    region_sub_tab["signif_status"] = ""
    # blue: significant pvalues for omnibus and univariate tests
    cond = np.where((region_sub_tab.JASS_PVAL < significance_treshold) & (region_sub_tab.UNIVARIATE_MIN_PVAL < significance_treshold))[0]
    region_sub_tab.loc[region_sub_tab.index[cond], "signif_status"] = "Both"

    # red: significant pvalues for omnibus test only
    cond = np.where((region_sub_tab.JASS_PVAL < significance_treshold ) & (region_sub_tab.UNIVARIATE_MIN_PVAL > significance_treshold))[0]
    region_sub_tab.loc[region_sub_tab.index[cond], "signif_status"] = "Joint"

    # green: significant pvalues for univariate test only
    cond = np.where((region_sub_tab.JASS_PVAL > significance_treshold) & (region_sub_tab.UNIVARIATE_MIN_PVAL < significance_treshold))[0]
    region_sub_tab.loc[region_sub_tab.index[cond], "signif_status"] = "Univariate"

    # grey: non significant pvalues
    cond = np.where((region_sub_tab.JASS_PVAL > significance_treshold) & (region_sub_tab.UNIVARIATE_MIN_PVAL > significance_treshold))[0]
    region_sub_tab.loc[region_sub_tab.index[cond], "signif_status"] = "None"

    return region_sub_tab

def get_region_summary(sum_stat_tab, phenotype_ids, significance_treshold=5*10**-8):

    # Select the most significant SNP for the joint test for each region
    region_sub_tab = sum_stat_tab.sort_values("JASS_PVAL").groupby("Region").first()#.reset_index()

    # add minimum univariate p-value
    univar = sum_stat_tab.groupby("Region").min().UNIVARIATE_MIN_PVAL
    region_sub_tab.loc[univar.index, "UNIVARIATE_MIN_PVAL"] = univar.values
    # Tag SNPs depending on which test is significant
    region_sub_tab.reset_index(inplace=True)
    region_sub_tab = add_signif_status_column(region_sub_tab, significance_treshold)

    #  reorder columns
    region_sub_tab = region_sub_tab[['Region', "MiddlePosition", "snp_ids","CHR", "position", "Ref_allele", "Alt_allele", "JASS_PVAL", "UNIVARIATE_MIN_PVAL", "signif_status"] + phenotype_ids]

    return region_sub_tab

def post_computation_filtering(worktable_chunk, significant_treshold = 5*10**-8):
    """
    Remove SNPs that seems aberrant: SNPs with a very low p-value that are isolated
     in their region

    :param worktable_chunk: pandas DataFrame worktable part
    :type worktable_chunk: pandas DataFrame
    :param significant_treshold: threshold at which a p-value is considered significant
    :type significant_treshold: float
    """

    def count_nearly_significant(rsnp):
        return((rsnp.JASS_PVAL < (significant_treshold*20)).sum())

    res = worktable_chunk.groupby("Region").apply(count_nearly_significant)

    # select region with only one SNP that is significant which is
    # suspect
    reg = res.loc[res==1].index

    for reg_aberant in reg:
        aberant_SNP = worktable_chunk.loc[worktable_chunk.Region==reg_aberant].sort_values("JASS_PVAL").index[0]
        worktable_chunk.drop(aberant_SNP, inplace=True)

    return worktable_chunk

def compute_pleiotropy_index(W,significance_treshold):

    N_significatif = (2.0 * spst.norm.sf(W.fillna(0, inplace=False).abs())<significance_treshold).sum(1)
    N_pheno = (~W.isnull()).sum(1)
    #pleiotropy index is not meaningful for too few phenotype
    S = N_significatif/N_pheno
    S.loc[N_pheno < 4] = np.nan
    return S

def create_worktable_file(
    phenotype_ids: List[str],
    init_file_path: str,
    project_hdf_path: str,
    remove_nan: bool,
    stat: str = "jass.models.stats:omnibus_stat",
    optim_na: bool = True,
    csv_file: str = None,
    chunk_size: int = 50,
    significance_treshold = 5*10**-8,
    post_filtering=True,
    **kwargs
):
    """
    Create a worktable file from an initial data table by specifying the
    selected phenotypes and the computation strategy

    :param phenotype_ids: the list of IDs for the phenotypes to select
    :type phenotype_ids: list
    :param init_file_path: path to the initial data table
    :type init_file_path: str
    :param project_hdf_path: path to the worktable file that will be produced
    :type project_hdf_path: str
    :param remove_nan: boolean to control the JOST computation strategy:
                      - if True any SNP which has a NaN value for one of its \
                        phenotypes will be removed.
                      - if False SNPs will be removed only if all phenotype \
                        values are NaN and JOST will be performed the "smart" \
                        way to compensate these missing values.
    :type remove_nan: bool
    :param optim_na: boolean to control if we use a smart gestion of z NA \
    values. Should always be set to True except for performance tests.
    :type optim_na: bool
    :param chunk_size: the size of the chunks of the initial data
                       to process together. default is 50 and should probably not
                       be touched for anything else than experimentation
    :type chunk_size: int
    :param significant_treshold: threshold at which a p-value is considered significant
    :type significant_treshold: float
    :
    """
    # read data by chunks to optimize memory usage
    # select only rows (SNPs) where there are no missing data
    how_dropna = "any" if remove_nan else "all"
    if os.path.exists(project_hdf_path):
        os.remove(project_hdf_path)
    hdf_work = HDFStore(project_hdf_path)

    # subset of phenotypes that have been selected
    phenolist = read_hdf(init_file_path, "PhenoList")
    phenolist = phenolist.loc[phenotype_ids]
    hdf_work.put(
       "PhenoList", phenolist
    )

    # subset of covariance matrix for the selected phenotypes
    cov = read_hdf(init_file_path, "COV")
    sub_cov = cov.loc[phenotype_ids, phenotype_ids]
    hdf_work.put(
        "COV", sub_cov, format="table", data_columns=True
    )  # Covariance matrix

    regions = read_hdf(init_file_path, "Regions").index.tolist()
    sum_stat_tab_min_itemsizes = {"snp_ids": 50, "Region": 10, "CHR": 5}
    #['Region', "MiddlePosition", "snp_ids","CHR", "position", "Ref_allele", "Alt_allele", "JASS_PVAL", "UNIVARIATE_MIN_PVAL", "signif_status"]
    region_sub_table_min_itemsizes = {"Region": 10, "index": 10, "CHR": 5, "snp_ids": 50, "signif_status":20}

    smart_na_computation = not (remove_nan)
    module_name, function_name = stat.split(":")
    stat_module = importlib.import_module(module_name)
    stat_fn = getattr(stat_module, function_name)

    stat_compute = choose_stat_function(smart_na_computation, optim_na, function_name, stat_fn, sub_cov, samp_size=phenolist['Effective_sample_size'], **kwargs)

    Nchunk = len(regions) // chunk_size + 1

    Nsnp_total = 0
    Nsnp_jassed = 0

    for chunk in range(Nchunk):
        binf = chunk * chunk_size
        bsup = (chunk+1) * chunk_size
        sum_stat_tab = read_hdf(init_file_path, 'SumStatTab', columns=[
        'Region', 'CHR', 'position', 'snp_ids', 'Ref_allele', 'Alt_allele', 'MiddlePosition'] + phenotype_ids,
        where='Region >= {0} and Region < {1}'.format(binf, bsup))
        print("Regions {0} to {1}\r".format(binf, bsup))

        # Remake row index unique: IMPORTANT for assignation with .loc at line 98
        sum_stat_tab.dropna(
            axis=0, subset=phenotype_ids, how=how_dropna, inplace=True
        )
        sum_stat_tab.reset_index(drop=True, inplace=True)

        if sum_stat_tab.shape[0]==0:
            print("No data available for region {0} to region {1}".format(binf, bsup))
            continue # skip region if no data are available

        N_pheno = len(phenotype_ids)
        Nsnp_total = Nsnp_total + sum_stat_tab.shape[0]

        if remove_nan or stat.split(":")[-1] != "omnibus_stat":
            sum_stat_tab['JASS_PVAL'] = stat_compute(sum_stat_tab[phenotype_ids])
        else:
            # Sort SumStatTab by missing patterns
            patterns_missing, frequent_pattern  = compute_frequent_missing_pattern(sum_stat_tab[phenotype_ids])

            # Apply the statistic computation by missing patterns
            for pattern in frequent_pattern:
                bool_serie = (patterns_missing == pattern)
                try:
                    sum_stat_tab.loc[bool_serie, "JASS_PVAL"] = stat_compute(
                        sum_stat_tab.loc[bool_serie, phenotype_ids]
                    )
                except ValueError:
                    print("worktable")

            sum_stat_tab.index = Index(patterns_missing)
            # Keep_only frequent_pattern
            sum_stat_tab = sum_stat_tab.loc[frequent_pattern]
            # drop pattern index :
            sum_stat_tab.reset_index(drop=True, inplace=True)

        Nsnp_jassed = Nsnp_jassed + sum_stat_tab.shape[0]
        sum_stat_tab.sort_values(by=["Region", "CHR"], inplace=True)

        sum_stat_tab["UNIVARIATE_MIN_PVAL"] = DataFrame(
                    2.0 * spst.norm.sf(sum_stat_tab[phenotype_ids].fillna(0, inplace=False).abs()),
                    index=sum_stat_tab.index,
                ).min(axis=1)

        sum_stat_tab["UNIVARIATE_MIN_QVAL"] =  sum_stat_tab["UNIVARIATE_MIN_PVAL"]*(1-np.isnan(sum_stat_tab[phenotype_ids]).astype(int)).sum(1)
        sum_stat_tab.loc[sum_stat_tab.UNIVARIATE_MIN_QVAL>1 , "UNIVARIATE_MIN_QVAL"] = 1

        #Computing pleiotropy
        sum_stat_tab["PLEIOTROPY_INDEX"] = compute_pleiotropy_index(sum_stat_tab[phenotype_ids], significance_treshold)

        sum_stat_tab = sum_stat_tab[
            ["Region", "CHR", "snp_ids", "position", 'Ref_allele', 'Alt_allele',  "MiddlePosition", "JASS_PVAL", "UNIVARIATE_MIN_PVAL", "UNIVARIATE_MIN_QVAL","PLEIOTROPY_INDEX" ]
            + phenotype_ids]

        if post_filtering:
            sum_stat_tab = post_computation_filtering(sum_stat_tab)

        hdf_work.append(
            "SumStatTab", sum_stat_tab, min_itemsize=sum_stat_tab_min_itemsizes
        )

        if csv_file is not None:
            with open(csv_file, 'a') as f:
                sum_stat_tab.to_csv(f, header=f.tell()==0)

        region_sub_table = get_region_summary(sum_stat_tab, phenotype_ids, significance_treshold=significance_treshold)

        hdf_work.append(
            "Regions",
            region_sub_table
            ,min_itemsize=region_sub_table_min_itemsizes
        )
    hdf_work.close()

    print("{1} SNPs treated on {0} SNPs".format(Nsnp_jassed, Nsnp_total))

    RegionSubTable = read_hdf(project_hdf_path, "Regions")

    pval_min = RegionSubTable["UNIVARIATE_MIN_PVAL"]
    jost_min = RegionSubTable["JASS_PVAL"]
    summaryTable = DataFrame(
        np.array(
            [
                [
                    sum((jost_min < significance_treshold) & (pval_min < significance_treshold)),
                    sum((jost_min < significance_treshold) & (pval_min > significance_treshold)),
                ],
                [
                    sum((jost_min > significance_treshold) & (pval_min < significance_treshold)),
                    sum((jost_min > significance_treshold) & (pval_min > significance_treshold)),
                ],
            ]
        )
    )
    summaryTable.columns = ["PhenoSignif", "NoPhenoSignif"]
    summaryTable.index = ["JOSTSignif", "NoJOSTSignif"]

    hdf_work = HDFStore(project_hdf_path)
    hdf_work.put(
        "summaryTable", summaryTable, format="table", data_columns=True
    )  # Summary Table (contigency table)
    hdf_work.close()


def compute_frequent_missing_pattern(sum_stat_tab):
    """
    Compute the frequency of missing pattern in the dataset
    determine the number of pattern that should be taken into account to
    cover 99 percent of the snps
    """
    N_pheno = sum_stat_tab.shape[1]
    patterns_missing = Series(np.dot((1- sum_stat_tab.isnull()), 10**np.arange((N_pheno-1), -1, -1)))
    pattern_frequency = patterns_missing.value_counts() / len(patterns_missing)
    n_pattern = pattern_frequency.shape[0]#(pattern_frequency.cumsum() < 0.99).sum() + 1
    #n_pattern = min(5000, n_pattern)
    print("Number of pattern {}".format(n_pattern))
    #print("Covering {}/1 of snps ".format(pattern_frequency.iloc[:n_pattern].sum()))
    frequent_pattern  = pattern_frequency.index.tolist() #[:n_pattern]
    return patterns_missing, frequent_pattern


def stringize_dataframe_region_chr(dataframe: DataFrame):
    """
    Reformat Region and Chromosome numbers as strings in a DataFrame.
    This will eventually be deprecated once the js code accepts integers instead
    of strings

    :param project_hdf_path: path to the worktable file
    :type project_hdf_path: str
    :param frame: name of the frame to be read
    :type frame: str
    :return: The dataframe with converted Region and CHR columns
    :rtype: pandas.DataFrame
    """
    dataframe["Region"] = dataframe["Region"].apply(lambda x: "Region" + str(x))
    dataframe["CHR"] = dataframe["CHR"].apply(lambda x: "chr" + str(x))
    return dataframe


def get_worktable_summary(project_hdf_path: str):
    """
    Read and return the summaryTable dataframe from a worktable file

    :param project_hdf_path: path to the worktable file
    :type project_hdf_path: str
    :return: The dataframe as a dictionary with one entry per line number
    :rtype: dict
    """
    summary_dataframe = read_hdf(project_hdf_path, "summaryTable")
    lines = {}
    for index, row in summary_dataframe.iterrows():
        lines[index] = {
            "PhenoSignif": int(row["PhenoSignif"]),
            "NoPhenoSignif": int(row["NoPhenoSignif"]),
        }
    return lines


def get_worktable_genomedata(project_hdf_path: str):
    """
    Read and return the Regions dataframe from a worktable file

    :param project_hdf_path: path to the worktable file
    :type project_hdf_path: str
    :return: The dataframe as a CSV formatted text
    :rtype: str
    """
    region_subtable = stringize_dataframe_region_chr(read_hdf(project_hdf_path, "Regions"))
    region_subtable.rename(index=str, columns={'JASS_PVAL':'JOSTmin'}, inplace=True)
    return region_subtable.to_csv(index=False)


def get_worktable_local_manhattan_data(
    project_hdf_path: str, chromosome: str, region: str
):
    """
    Read and return the SumStatTab dataframe from a worktable file
    for a given chromosome and region for the Manhattan plot

    :param project_hdf_path: path to the worktable file
    :type project_hdf_path: str
    :param chromosome: chromosome number in "string form", e.g. chr12
    :type chromosome: str
    :param region: region number in "string form", e.g. Region213
    :type region: str
    :return: The dataframe subset corresponding to the chromosome and region, as a CSV formatted text
    :rtype: str
    """
    region_int = region[6:]
    chromosome_int = chromosome[3:]
    dataframe = read_hdf(project_hdf_path, "SumStatTab",
                            columns=["Region", "CHR", "position", "snp_ids", "JASS_PVAL"],
                            where=['Region='+str(region_int), 'CHR='+str(chromosome_int)])
    dataframe = stringize_dataframe_region_chr(dataframe)
    dataframe = dataframe.sort_values("position")
    return dataframe.to_csv(index=False)


def get_worktable_local_heatmap_data(
    project_hdf_path: str, chromosome: str, region: str
):
    """
    Read and return the SumStatTab dataframe from a worktable file
    for a given chromosome and region for the Heatmap plot

    :param project_hdf_path: path to the worktable file
    :type project_hdf_path: str
    :param chromosome: chromosome number in "string form", e.g. chr12
    :type chromosome: str
    :param region: region number in "string form", e.g. Region213
    :type region: str
    :return: The dataframe subset corresponding to the chromosome and region, \
             pivoted and as a CSV formatted text
    :rtype: str
    """
    region_int = region[6:]
    chromosome_int = chromosome[3:]
    dataframe = read_hdf(project_hdf_path, "SumStatTab",
                            where=['Region='+str(region_int), 'CHR='+str(chromosome_int)])
    dataframe = stringize_dataframe_region_chr(dataframe)
    dataframe = dataframe.sort_values("position")
    dataframe.drop( ["Region", "CHR", "position", "JASS_PVAL", "MiddlePosition"],
        axis=1,
        inplace=True,
    )
    dataframe.rename(columns={"snp_ids": "ID"}, inplace=True)
    column_order = list(dataframe.ID)
    pivoted_dataframe = dataframe.pivot_table(columns="ID")
    pivoted_dataframe = pivoted_dataframe.reindex_axis(column_order, axis=1)
    # TODO rework the selection to return 5000 snps in total, around the
    # region ;)
    return pivoted_dataframe.to_csv(index_label="ID")
