# Apache License Version 2.0
# Copyright 2022 Xin Huang
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import allel
import gzip
import os
import numpy as np
import pandas as pd
from multiprocessing import Process, Queue
from sstar.utils import read_data, py2round, read_mapped_region_file, cal_match_pct

#@profile
def cal_pvalue(vcf, ref_ind_file, tgt_ind_file, src_ind_file, anc_allele_file, output, thread, score_file, ref_match_pct_file, mapped_region_file, low_memory, mapped_len_esp, len_esp, var_esp, sfs_esp):
    """
    Description:
        Calculate p-values for S* haplotypes in the target population with source genomes.

    Arguments:
        vcf str: Name of the VCF file containing genotypes.
        src_vcf str: Name of the VCF file containing genotypes from source populations.
        ref_ind_file str: Name of the file containing sample information from reference populations.
        tgt_ind_file str: Name of the file containing sample information from target populations.
        src_ind_file str: Name of the file containing sample information from source populations.
        anc_allele_file str: Name of the file containing ancestral allele information.
        output str: Name of the output file.
        thread int: Number of threads.
        score_file str: Name of the file containing S* scores calculated by `s-star score`.
        ref_match_pct_file str: Names of the files containing match percents in reference populations calculated by `s-star rmatch`.
        mapped_region_file str: Name of the BED file containing mapped regions.
        mapped_len_esp float: Increment of the length of the mapped region.
        len_esp float: Increment of the length of the haplotype.
        var_esp float: Increment of the number of derived alleles on the haplotype.
        sfs_esp float: Increment of mean site frequency spectrum.
    """

    ref_data, ref_samples, tgt_data, tgt_samples, src_data, src_samples = read_data(vcf, ref_ind_file, tgt_ind_file, src_ind_file, anc_allele_file)

    res = []
    chr_names = ref_data.keys()

    mapped_intervals = read_mapped_region_file(mapped_region_file)
    data, windows, samples = _read_score_file(score_file, chr_names, tgt_samples)
    sample_size = len(samples)

    header = 'chrom\tstart\tend\tsample\tp-value\t'
    header += 'src_sample\thap_index\tS*_start\tS*_end\tS*_SNP_num\t'
    header += "hap_dSNV_num\thap_len\thap_mapped_len\thap_match_num\thap_tot_num\thap_dSNP_per_site_num\thap_S*_match(%)\thap_num_match_ref"

    # Read match percents in reference populations from a file
    # Use whole-genome match percents as the null distributions

    if low_memory:
        try:
           ref_match_pct = pd.read_csv(ref_match_pct_file, compression="gzip", sep="\t")
        except:
           ref_match_pct = pd.read_csv(ref_match_pct_file, sep="\t")
        query_ref_match_pct = _query_ref_match_pct_pandas
    else:
        ref_match_pct = _read_ref_match_pct_file(ref_match_pct_file)
        query_ref_match_pct = _query_ref_match_pct_naive

    #for s in samples[0:1]:
    #    i = samples.index(s)
    #    res = _cal_pvalue_ind(data[s], i, mapped_intervals, tgt_data, src_data, src_samples, ref_match_pct, sample_size, query_ref_match_pct, mapped_len_esp, len_esp, var_esp, sfs_esp)

    if thread is None: thread = min(os.cpu_count()-1, sample_size)    
    res = _cal_tgt_match_pct_manager(data, mapped_intervals, samples, tgt_samples, src_samples, tgt_data, src_data, ref_match_pct, sample_size, query_ref_match_pct, thread, mapped_len_esp, len_esp, var_esp, sfs_esp)

    with open(output, 'w') as o:
        o.write(header+"\n")
        o.write("\n".join(res)+"\n")

#@profile
def _read_score_file(score_file, chr_names, tgt_samples):
    """
    Description:
        Helper function for reading the file generated by `sstar score`.

    Arguments:
        score_file str: Name of the file containing S* scores generated by `sstar score`.
        chr_names list: List containing names of chromosomes for analysis.
        tgt_samples list: List containing names of samples from the target population for analysis.

    Returns:
        data dict: Dictionary containing S* for analysis.
        windows dict: Dictionary containing windows for analysis.
        header str: Header from the file generated by `sstar score`.
        samples list: List containing names of samples in the target population for analysis.
    """

    data = dict()
    windows = dict()
    for c in chr_names:
        windows[c] = []
    samples = []
    with open(score_file, 'r') as f:
        header = f.readline().rstrip()
        for line in f.readlines():
            line = line.rstrip()
            elements = line.split("\t")
            chr_name = elements[0]
            win_start = elements[1]
            win_end = elements[2]
            sample = elements[3]
            if sample not in tgt_samples: continue
            if elements[6] == 'NA': continue
            if sample not in data.keys(): 
                data[sample] = []
                samples.append(sample)
            data[sample].append(line)
            star_snps = elements[-1].split(",")
            windows[c].append((int(win_start), int(win_end)))
            windows[c].append((int(star_snps[0]), int(star_snps[-1])))

    return data, windows, samples

#@profile
def _read_ref_match_pct_file(ref_match_pct_file):
    """
    Description:
        Helper function for reading match percents from the reference population.

    Arguments:
        ref_match_pct_file str: Name of the file containing match percents from the reference population.

    Returns:
        ref_match_pct dict: Dictionary containing match percents from the reference population.
    """

    f = gzip.open(ref_match_pct_file, 'rt')
    try:
        f.readline()
    except:
        f.close()
        f = open(ref_match_pct_file, 'r')
        f.readline()

    ref_match_pct = dict()
        
    for line in f.readlines():
        elements = line.rstrip().split("\t")
        count = int(elements[0])
        mapped_bases_bin = int(elements[1])
        hap_len = int(elements[2])
        mh_sites = int(elements[3])
        tot_sites = int(elements[4])
        sfs = float(elements[5])
        match = float(elements[6])
        if mapped_bases_bin not in ref_match_pct.keys(): ref_match_pct[mapped_bases_bin] = dict()
        if hap_len not in ref_match_pct[mapped_bases_bin].keys(): ref_match_pct[mapped_bases_bin][hap_len] = dict()
        if mh_sites not in ref_match_pct[mapped_bases_bin][hap_len].keys(): ref_match_pct[mapped_bases_bin][hap_len][mh_sites] = dict()
        if sfs not in ref_match_pct[mapped_bases_bin][hap_len][mh_sites].keys(): 
            ref_match_pct[mapped_bases_bin][hap_len][mh_sites][sfs] = dict()
            ref_match_pct[mapped_bases_bin][hap_len][mh_sites][sfs]['count'] = []
            ref_match_pct[mapped_bases_bin][hap_len][mh_sites][sfs]['match_pct'] = []
        ref_match_pct[mapped_bases_bin][hap_len][mh_sites][sfs]['count'].append(count)
        ref_match_pct[mapped_bases_bin][hap_len][mh_sites][sfs]['match_pct'].append(match / tot_sites)

    f.close()

    return ref_match_pct

def _cal_tgt_match_pct_manager(data, mapped_intervals, samples, tgt_samples, src_samples, tgt_data, src_data, ref_match_pct, sample_size, query_ref_match_pct, thread, mapped_len_esp, len_esp, var_esp, sfs_esp):
    """
    Description:
        Manager function to calculate match percents in target populations using multiprocessing.

    Arguments:
        data dict: Lines from the output file created by `sstar score`.
        mapped_intervals dict: Dictionary of tuples containing mapped regions across the genome.
        sample list: Sample information for individuals needed to be estimated match percents.
        tgt_samples list: Sample information from target populations.
        src_samples list: Sample information from source populations.
        tgt_data dict: Genotype data from target populations.
        src_data dict: Genotype data from source populations.
        ref_match_pct dict: Match percents calculated from reference populations.
        sample_size int: Number of individuals analyzed.
        query_ref_match_pct func: Function used to query match percentage from reference populations.
        thread int: Number of threads.
        mapped_len_esp float: Increment of the length of the mapped region.
        len_esp float: Increment of the length of the haplotype.
        var_esp float: Increment of the number of derived alleles on the haplotype.
        sfs_esp float: Increment of mean site frequency spectrum.

    Returns:
        res list: Match percents for target populations.
    """

    try:
        from pytest_cov.embed import cleanup_on_sigterm
    except ImportError:
        pass
    else:
        cleanup_on_sigterm()

    res = []
    in_queue, out_queue = Queue(), Queue()
    workers = [Process(target=_cal_tgt_match_pct_worker, args=(in_queue, out_queue, mapped_intervals, tgt_data, src_data, src_samples, ref_match_pct, len(tgt_samples), query_ref_match_pct, mapped_len_esp, len_esp, var_esp, sfs_esp)) for ii in range(thread)]

    for t in samples:
        index = tgt_samples.index(t)
        in_queue.put((index, data[t]))
    
    try:
        for worker in workers:
            worker.start()
        for s in range(sample_size):
            item = out_queue.get()
            if item != '': res.append(item)
        for worker in workers:
            worker.terminate()
    finally:
        for worker in workers:
            worker.join()

    return res

def _cal_tgt_match_pct_worker(in_queue, out_queue, mapped_intervals, tgt_data, src_data, src_samples, ref_match_pct, sample_size, query_ref_match_pct, mapped_len_esp, len_esp, var_esp, sfs_esp):
    """
    Description:
        Worker function to calculate match percents in target populations.

    Arguments:
        in_queue multiprocessing.Queue: multiprocessing.Queue instance to receive parameters from the manager.
        out_queue multiprocessing.Queue: multiprocessing.Queue instance to send results back to the manager.
        mapped_intervals dict: Dictionary of tuples containing mapped regions across the genome.
        tgt_data dict: Genotype data from target populations.
        src_data dict: Genotype data from source populations.
        src_samples list: List containing sample information for source populations.
        ref_match_pct dict: Match percents in reference populations as the null distribution.
        sample_size int: Number of individuals analyzed.
        query_ref_match_pct func: Function used to query match percentages from reference popualtions.
        mapped_len_esp float: Increment of the length of the mapped region.
        len_esp float: Increment of the length of the haplotype.
        var_esp float: Increment of the number of derived alleles on the haplotype.
        sfs_esp float: Increment of mean site frequency spectrum.
    """

    while True:
        index, data = in_queue.get()
        res = _cal_pvalue_ind(data, index, mapped_intervals, tgt_data, src_data, src_samples, ref_match_pct, sample_size, query_ref_match_pct, mapped_len_esp, len_esp, var_esp, sfs_esp)
        out_queue.put("\n".join(res))

#@profile
def _get_ssnps_range(chr_name, data, ind_index, hap_index, win_start, win_end, s_star_snps):
    """
    Description:
        Helper function to obtain the range of a haplotype containing S* SNPs.
        If the haplotype contains less than two S* SNPs, it will be ignored.

    Arguments:
        chr_name str: Name of the chromosome.
        data dict: Dictionary containing genotype data and position information.
        ind_index int: Index of the individual carrying S* SNPs.
        hap_index int: Index of the haplotype carrying S* SNPs.
        win_start int: Start position of the local window containing S* SNPs.
        wind_end int: End position of the local window containing S* SNPs.
        s_star_snps list: List containing positions of S* SNPs.

    Returns:
        hap_pos_min int: Start position of the haplotype.
        hap_pos_max int: End position of the haplotype.
    """
    gt = data[chr_name]['GT']
    pos = data[chr_name]['POS']

    sub_snps = np.where((pos>=win_start) & (pos<=win_end))[0]
    sub_gt = gt[sub_snps][:,ind_index]
    sub_pos = pos[sub_snps]

    hap = sub_gt[:,hap_index]
    s_star_snps_pos = [int(s) for s in s_star_snps]
    index = np.in1d(sub_pos, s_star_snps_pos)
    hap_num = np.sum(hap[index])
    if hap_num < 2: 
        hap_pos_max = 'NA'
        hap_pos_min = 'NA'
    else:
        hap_pos_max = int(np.array(s_star_snps_pos)[np.equal(hap[index],1)][-1])
        hap_pos_min = int(np.array(s_star_snps_pos)[np.equal(hap[index],1)][0])

    return hap_pos_min, hap_pos_max, hap_num

#@profile
def _query_ref_match_pct_naive(ref_match_pct, hap_stats, mapped_len_esp, len_esp, var_esp, sfs_esp):
    """
    Description:
        Helper function for querying match percents from the reference population 
        consistent with statistics from a given haplotype using Python dictionary.

    Arguments:
        ref_match_pct dict: Dictionary containing match percents from the reference population.
        hap_stats list: List containing statistics from a haplotype.
        mapped_len_esp float: Increment of the length of the mapped region.
        len_esp float: Increment of the length of the haplotype.
        var_esp float: Increment of the number of derived alleles on the haplotype.
        sfs_esp float: Increment of mean site frequency spectrum.

    Returns:
        res list: List containing match percents from the reference population.
    """

    res = []

    mapped_len = hap_stats[2]
    hap_len = hap_stats[1]
    hap_site = hap_stats[-3]
    hap_var = hap_stats[0]
    sfs = hap_stats[-2]

    if (mapped_len == 'NA') or (hap_len == 'NA') or (hap_var == 'NA') or (sfs == 'NA'): return res
    mapped_len_start = mapped_len - mapped_len_esp
    mapped_len_end = mapped_len + mapped_len_esp
    len_start = hap_len - len_esp
    len_end = hap_len + len_esp
    var_start = hap_var - var_esp
    var_end = hap_var + var_esp
    sfs_start = sfs - sfs_esp
    sfs_end = sfs + sfs_esp

    for mapped_len in ref_match_pct.keys():
        if (mapped_len >= mapped_len_start) and (mapped_len <= mapped_len_end):
            for hap_len in ref_match_pct[mapped_len].keys():
                if (hap_len >= len_start) and (hap_len <= len_end):
                    for hap_var in ref_match_pct[mapped_len][hap_len].keys():
                        if (hap_var >= var_start) and (hap_var <= var_end):
                            for sfs in ref_match_pct[mapped_len][hap_len][hap_var].keys():
                                if (sfs >= sfs_start) and (sfs <= sfs_end):
                                    for e in zip(ref_match_pct[mapped_len][hap_len][hap_var][sfs]['match_pct'], ref_match_pct[mapped_len][hap_len][hap_var][sfs]['count']):
                                        res += [e[0]] * e[1]

    return res

#@profile
def _query_ref_match_pct_pandas(df, hap_stats, mapped_len_esp, len_esp, var_esp, sfs_esp):
    """
    Description:
        Helper function for querying match percents from the reference population 
        consistent with statistics from a given haplotype using the polars package.

    Arguments:
        ref_match_pct dict: Dictionary containing match percents from the reference population.
        hap_stats list: List containing statistics from a haplotype.
        mapped_len_esp float: Increment of the length of the mapped region.
        len_esp float: Increment of the length of the haplotype.
        var_esp float: Increment of the number of derived alleles on the haplotype.
        sfs_esp float: Increment of mean site frequency spectrum.

    Returns:
        res list: List containing match percents from the reference population.
    """

    res = []

    mapped_len = hap_stats[2]
    hap_len = hap_stats[1]
    hap_site = hap_stats[-3]
    hap_var = hap_stats[0]
    sfs = hap_stats[-2]

    if (mapped_len == 'NA') or (hap_len == 'NA') or (hap_var == 'NA') or (sfs == 'NA'): return res
    mapped_len_start = mapped_len - mapped_len_esp
    mapped_len_end = mapped_len + mapped_len_esp
    len_start = hap_len - len_esp
    len_end = hap_len + len_esp
    var_start = hap_var - var_esp
    var_end = hap_var + var_esp
    sfs_start = sfs - sfs_esp
    sfs_end = sfs + sfs_esp

    counts = df[(df['hap_len'] >= len_start) & (df['hap_len'] <= len_end) & (df['hap_mapped_len_bin'] >= mapped_len_start) & (df['hap_mapped_len_bin'] <= mapped_len_end) & (df['hap_match_num'] >= var_start)& (df['hap_match_num'] <= var_end) & (df['hap_dSNP_per_site_num'] >= sfs_start) & (df['hap_dSNP_per_site_num'] <= sfs_end)]

    matches = counts['match_pct'].values
    counts = counts['count'].values
    for i in range(len(counts)):
        res += [matches[i]] * counts[i]

    return res

#@profile
def _cal_pvalue(ref_match_pct, hap_match_pct):
    """
    Description:
        Helper function to calculate a p-value for a given match percent.

    Arguments:
        ref_match_pct list: Match percents from reference populations as the null distribution.
        hap_match_pct float: Match percent of a haplotype in a window.

    Returns:
        pvalue float: P-value for the significance of the match percent.
    """

    if (ref_match_pct is None) or (len(ref_match_pct) == 0): return 'NA'
    if hap_match_pct == 0: return 1.0
    if hap_match_pct == 'NA': return 'NA'
    count = len([r for r in ref_match_pct if hap_match_pct <= round(r, 6)])
    pvalue = count / len(ref_match_pct)
    pvalue = round(pvalue, 6)

    return pvalue

#@profile
def _cal_pvalue_ind(data, tgt_ind_index, mapped_intervals, tgt_data, src_data, src_samples, ref_match_pct, sample_size, query_ref_match_pct, mapped_len_esp, len_esp, var_esp, sfs_esp):
    """
    Description:
        Helper function for calculating p-values in individuals.

    Arguments:
        data dict: Dictionary containing S* for analysis.
        tgt_ind_index int: Index of the target individual for analysis.
        mapped_intervals dict: Dictionary of tuples containing mapped regions across the genome.
        tgt_data dict: Genotype data from target populations.
        src_data dict: Genotype data from source populations.
        src_samples list: List of samples from source populations.
        ref_match_pct dict: Dictionary containing match percents from the reference population.
        sample_size int: Number of individuals analyzed.
        query_ref_match_pct func: Function used to query match percentages from reference popualtions.
        mapped_len_esp float: Increment of the length of the mapped region.
        len_esp float: Increment of the length of the haplotype.
        var_esp float: Increment of the number of derived alleles on the haplotype.
        sfs_esp float: Increment of mean site frequency spectrum.

    Returns:
        res list: List containing estimated p-values and other statistics.
    """
    res = []
    for line in data:
        elements = line.split("\t")
        chr_name = elements[0]
        win_start, win_end = elements[1], elements[2]
        sample = elements[3]

        s_star_snps = elements[-1].split(",")
        s_start, s_end = s_star_snps[0], s_star_snps[-1]
        key1 = win_start+'-'+win_end
        key2 = s_start+'-'+s_end

        for src_ind_index in range(len(src_samples)):
            # Results using SNPs within S* SNPs
            src_sample = src_samples[src_ind_index]
            hap1_s_start, hap1_s_end, hap1_s_num = _get_ssnps_range(chr_name, tgt_data, tgt_ind_index, 0, int(win_start), int(win_end), s_star_snps)
            hap2_s_start, hap2_s_end, hap2_s_num = _get_ssnps_range(chr_name, tgt_data, tgt_ind_index, 1, int(win_start), int(win_end), s_star_snps)
            hap1_star_res = cal_match_pct(chr_name, mapped_intervals, tgt_data, src_data, tgt_ind_index, src_ind_index, 0, hap1_s_start, hap1_s_end, sample_size)
            hap2_star_res = cal_match_pct(chr_name, mapped_intervals, tgt_data, src_data, tgt_ind_index, src_ind_index, 1, hap2_s_start, hap2_s_end, sample_size)
            
            hap1_star_match_pct = hap1_star_res[-1]
            hap2_star_match_pct = hap2_star_res[-1]

            ref_match_pct3 = query_ref_match_pct(ref_match_pct, hap1_star_res, mapped_len_esp, len_esp, var_esp, sfs_esp)
            ref_match_pct4 = query_ref_match_pct(ref_match_pct, hap2_star_res, mapped_len_esp, len_esp, var_esp, sfs_esp)

            hap1_star_match_pvalue = _cal_pvalue(ref_match_pct3, hap1_star_match_pct)
            hap2_star_match_pvalue = _cal_pvalue(ref_match_pct4, hap2_star_match_pct)

            hap1_num_match_ref = 'NA'
            hap2_num_match_ref = 'NA'
            if ref_match_pct3 is not None: hap1_num_match_ref = len(ref_match_pct3)
            if ref_match_pct4 is not None: hap2_num_match_ref = len(ref_match_pct4)

            line1 = f'{chr_name}\t{win_start}\t{win_end}\t{sample}\t{hap1_star_match_pvalue}\t'
            line1 += f'{src_sample}\t1\t{hap1_s_start}\t{hap1_s_end}\t{hap1_s_num}\t'
            line1 += "\t".join([str(r) for r in hap1_star_res]) + f"\t{hap1_num_match_ref}"
            line2 = f'{chr_name}\t{win_start}\t{win_end}\t{sample}\t{hap2_star_match_pvalue}\t'
            line2 += f'{src_sample}\t2\t{hap2_s_start}\t{hap2_s_end}\t{hap2_s_num}\t'
            line2 += "\t".join([str(r) for r in hap2_star_res]) + f"\t{hap2_num_match_ref}"
            res.append(line1)
            res.append(line2)
    return res
