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

@author: vguillem, hjulienne, hmenager
"""
import os
import glob
import logging
from pandas import HDFStore, DataFrame, read_csv, concat, options

# create (or open) an hdf5 file and opens in append mode
import numpy as np
import tables
import warnings
from functools import reduce

options.mode.chained_assignment = None
warnings.filterwarnings("ignore", category=tables.NaturalNameWarning)

def get_gwasname(file_name):
    return "_".join(os.path.basename(file_name).split("_")[0:3])


def check_if_SNP_unique(z_gwas_chrom):
    tg = z_gwas_chrom.groupby(["snp_ids"]).count()["CHR"].reset_index()
    if len(tg.loc[tg["CHR"] > 1, "CHR"]) > 0:
        msg = "Bad input : SNP {0} have several positions or reference allele".format(
            " ".join(tg.loc[tg["CHR"] > 0, "snp_ids"])
        )
        raise IOError(msg)

def get_gwas_dict(input_data_path):
    gwas_dict = {}
    #retrieve all files corresponding to glob patterns
    files_by_pattern = list(map(glob.glob, input_data_path.split(";")))
    all_path = reduce(lambda x,y:x+y, files_by_pattern)
    print(all_path)
    for x in all_path:
        gwas_name = get_gwasname(x)

        if gwas_name in gwas_dict:
            gwas_dict[gwas_name].append(x)
        else:
            gwas_dict[gwas_name] = [x]
    return gwas_dict

def create_pheno_summary(description):
    pheno_summary_data = description[
        [
            "Consortium",
            "Outcome",
            "FullName",
            "Type",
            "Reference",
            "ReferenceLink",
            "dataLink",
            "internalDataLink",
            "Nsample",
            "Ncase",
            "Ncontrol"
        ]
    ]

    pheno_summary_data["ID"] =  "z_" + pheno_summary_data.Consortium.str.upper() + "_" + pheno_summary_data.Outcome.str.upper()

    is_quantitatif = pheno_summary_data['Ncase'].isnull()
    pheno_summary_data["Effective_sample_size"] = np.nan
    pheno_summary_data.loc[~is_quantitatif,"Effective_sample_size"] = (pheno_summary_data.loc[~is_quantitatif, "Ncase"]*pheno_summary_data.loc[~is_quantitatif, "Ncontrol"]) / pheno_summary_data.loc[~is_quantitatif, "Nsample"]

    pheno_summary_data.loc[is_quantitatif,"Effective_sample_size"] = pheno_summary_data.loc[is_quantitatif, "Nsample"]

    # reorder columns in the dataframe
    pheno_summary_data = pheno_summary_data[
        [
            "ID",
            "Consortium",
            "Outcome",
            "FullName",
            "Type",
            "Reference",
            "ReferenceLink",
            "dataLink",
            "internalDataLink",
            "Nsample",
            "Ncase",
            "Ncontrol",
            "Effective_sample_size"
        ]
    ]
    pheno_summary_data.index = pheno_summary_data["ID"]



    return pheno_summary_data

def format_chr_gwas(gwas_file_chri, chrom, study_name, regions_bychr):

    z_gwas = read_csv(
        gwas_file_chri, sep="\t", usecols=[0, 1, 2, 3, 4], memory_map=True
    )
    z_gwas.columns = [
        "snp_ids",
        "position",
        "Ref_allele",
        "Alt_allele",
        "z_score",
    ]

    # CAUTION: handle duplicates (whenever there are multiple values
    # for a single SNP)
    z_gwas.drop_duplicates(subset=["snp_ids"], keep="first", inplace=True)

    # Add chromosome and study name columns
    z_gwas["CHR"] = chrom
    z_gwas["study_name"] = study_name

    # Add region and middle position columns
    z_gwas["Region"] = 0
    z_gwas["MiddlePosition"] = 0.0

    for region_index, region_row in regions_bychr.get_group(
        "chr%d" % chrom
    ).iterrows():
        left = region_row["start"]
        right = region_row["stop"]
        ind = (z_gwas["position"] >= left) & (z_gwas["position"] <= right)
        (z_gwas.loc[ind, "Region"]) = np.int(region_index + 1)
        (z_gwas.loc[ind, "MiddlePosition"]) = (left + right) / 2
    return(z_gwas)

def compute_covariance_zscore(init_table_path):
    print("## Compute covariance ##")
    sum_stat_jost_tab = pd.read_hdf(init_file_path, 'SumStatTab', where='Region >= {0} and Region < {1}'.format(0, 3))
    trait = [i for i in sum_stat_jost_tab.columns if i[:2]=="z_"]

    NSNP_matrix = pd.DataFrame(index=trait, columns=trait)
    cov_matrix = pd.DataFrame(index=trait, columns=trait)

    cov_matrix.fillna(0, inplace=True)
    NSNP_matrix.fillna(0, inplace=True)

    bi = range(0,1751,50)
    start = timeit.timeit()
    n_len = len(bi)-1

    for i in range(n_len):
        binf = bi[i]
        bsup = bi[(i+1)]
        sum_stat_jost_tab = pd.read_hdf(init_file_path, 'SumStatTab', where='Region >= {0} and Region < {1}'.format(binf, bsup))

        print("Regions {0} to {1}\r".format(binf, bsup))
        j = 0
        for tr1 in trait:
            for tr2 in trait[j:]:
                cc = sum_stat_jost_tab[[tr1,tr2]].dropna()
                cc = cc.loc[cc.max(1) < 4]

                cov_matrix.loc[tr1,tr2] += cc.iloc[:,0].dot(cc.iloc[:,1])
                NSNP_matrix.loc[tr1,tr2] += cc.shape[0]

                cov_matrix.loc[tr2,tr1] += cc.iloc[:,0].dot(cc.iloc[:,1])
                NSNP_matrix.loc[tr2,tr1] += cc.shape[0]
            j=j+1
    end = timeit.timeit()
    print(end-start)
    #(cov_matrix/NSNP_matrix).to_csv("Covariance_on_Zscores.csv", sep="\t")
    hdf_init = pd.HDFStore(init_table_path)
    hdf_init.put("COV", (cov_matrix/NSNP_matrix), format="table", data_columns=True)
    hdf_init.close()

def create_inittable_file(
    input_data_path: str,
    init_covariance_path,
    regions_map_path: str,
    description_file_path: str,
    init_table_path: str,
):
    # Read region file
    regions = read_csv(regions_map_path, sep='\s+', memory_map=True)
    # Create HDFStore
    if os.path.exists(init_table_path):
        os.remove(init_table_path)
    hdf_init = HDFStore(init_table_path)

    # Read covariance file
    if init_covariance_path!=None:
        covariance = read_csv(init_covariance_path, sep="\t", index_col=0, memory_map=True)
        compute_covariance = False
    else:
        compute_covariance = True

    # Read description file
    description = read_csv(description_file_path, sep="\t")
    gwas_dict = get_gwas_dict(input_data_path)
    print(gwas_dict)
    phenotypes_list = []
    for meta_index, meta_row in description.iterrows():
        # GWAS loop
        # create an ID from the GWAS consortia + outcome columns
        phenotypes_list.append(
            "z_" + meta_row["Consortium"].upper() + "_" + meta_row["Outcome"].upper()
        )

    pheno_select = list(phenotypes_list) & gwas_dict.keys() #&covariance.columns
    pheno_summary_data = create_pheno_summary(description)

    # select only phenotypes for which there is a covariance
    pheno_list = pheno_summary_data.loc[pheno_select, :]
    if compute_covariance==False:
        COV = covariance.loc[pheno_select, pheno_select]
        hdf_init.put("COV", COV, format="table", data_columns=True)

    which_cols = [
        "Region",
        "CHR",
        "MiddlePosition",
        "position",
        "snp_ids",
        "Ref_allele",
        "Alt_allele",
    ]

    which_cols.extend(list(pheno_select))
    hdf_init.put("PhenoList", pheno_list, format="table", data_columns=True)
    hdf_init.put("Regions", regions, format="table", data_columns=True)
    sum_stat_tab_min_itemsizes = {"snp_ids": 50}

    regions_bychr = regions.groupby("chr")

    for chrom in range(1, 23):  # Chromosome loop
        print(chrom)
        z_gwas_long_chrom = DataFrame()
        logging.debug("Processing chromosome %d" % chrom)
        for meta_index, meta_row in description.iterrows():  # GWAS loop
            # create an ID from the GWAS consortia + outcome columns
            study_name = (
                "z_"
                + meta_row["Consortium"].upper()
                + "_"
                + meta_row["Outcome"].upper()
            )

            print(study_name)
            logging.debug("Processing study %s" % study_name)
            logging.info("Processing study %s" % study_name)
            try:
                gwas_file_chri = os.path.join(
                    os.path.dirname(input_data_path),
                    study_name + "_chr" + str(chrom) + ".txt",
                )
                if gwas_file_chri not in gwas_dict[study_name]:
                    logging.warning(
                        "file %s missing for chromosome no.%s / GWAS %s"
                        % (gwas_file_chri, chrom, study_name)
                    )
                    continue
            except KeyError:
                logging.warning("files missing for GWAS %s" % (study_name))
                # As it is the pipeline can't handle missing files or empty files
                continue

            # Load summary statistics file for CHR i
            z_gwas = format_chr_gwas(gwas_file_chri, chrom, study_name, regions_bychr)
            z_gwas_long_chrom = concat(
                [
                    z_gwas_long_chrom,
                    z_gwas[
                        [
                            "snp_ids",
                            "position",
                            "Ref_allele",
                            "Alt_allele",
                            "Region",
                            "MiddlePosition",
                            "CHR",
                            "study_name",
                            "z_score",
                        ]
                    ],
                ],
                axis=0,
            )

        if len(z_gwas_long_chrom) > 0:
            z_gwas_chrom = z_gwas_long_chrom.pivot_table(
                index=[
                    "snp_ids",
                    "position",
                    "Ref_allele",
                    "Alt_allele",
                    "CHR",
                    "Region",
                    "MiddlePosition",
                ],
                columns="study_name",
                values="z_score",
            )
            z_gwas_chrom.reset_index(inplace=True)

            # Check if only one Reference_allele and position by SNP
            check_if_SNP_unique(z_gwas_chrom)

            sum_stat_tab = z_gwas_chrom[which_cols]
            hdf_init.append(
                "SumStatTab",
                sum_stat_tab,
                min_itemsize=sum_stat_tab_min_itemsizes,
                data_columns=["Region"],
            )
    hdf_init.close()

    if compute_covariance:
        print("Compute Covariance from Zscores")
        compute_covariance_zscore(hdf_init = init_table_path)
