import numpy as np

import pandas as pd

import torch

from scipy import interpolate

from .utils import load_dict

import pandas as pd

import numpy as np

from copy import copy

def find_closest_x(df_pred, lower_bound=0, upper_bound=200, tolerance=0.01):

    upper = 1e-2

    lower = 1e-3

    while lower_bound <= upper_bound:

        mid = (lower_bound + upper_bound) / 2

        #result = len(np.where(df_pred.P_weighted.values * mid < 2e-4)[0]) / len(np.where(df_pred.P.values < 2e-4)[0])

        res1 = len(np.where((df_pred.P_weighted.values * mid < upper) & (df_pred.P_weighted.values * mid > lower))[0])

        res2 = len(np.where((df_pred.P.values < upper) & (df_pred.P.values > lower))[0])

        result = res1/res2

        if abs(result - 1) < tolerance:

            return mid

        elif result > 1:

            lower_bound = mid + tolerance

        else:

            upper_bound = mid - tolerance

    return mid

def get_clumps_gold_label(data_path, gold_label_gwas, t_p = 5e-8, no_hla = False, column = 'P', snp2ld_snps = None):

    snp2ld_snps_with_hla = load_dict(data_path + 'ld_score/ukb_white_ld_10MB.pkl')

    snp2ld_snps_no_hla = load_dict(data_path + 'ld_score/ukb_white_ld_10MB_no_hla.pkl')

    if not snp2ld_snps:

        if no_hla:

            snp2ld_snps = snp2ld_snps_no_hla

        else:

            snp2ld_snps = snp2ld_snps_with_hla

    clumps = []

    snps_in_clumps = []

    snp_hits = gold_label_gwas[gold_label_gwas[column] < t_p].sort_values(column).SNP.values

    for snp in snp_hits:

        if snp in snps_in_clumps:

            ## already in existing clumps => not create a new clump

            pass

        else:

            if snp in snp2ld_snps:

                # ld block

                clumps.append([snp] + snp2ld_snps[snp])

                snps_in_clumps += snp2ld_snps[snp]

                snps_in_clumps += [snp]

            else:

                # no other SNPs tagged

                clumps.append([snp])

                snps_in_clumps += [snp]

    return clumps

def get_meta_clumps(clumps, data_path):

    snp2cm = dict(pd.read_csv(data_path + 'misc_data/ukb_white_with_cm.bim', sep = '\t', header = None)[[1, 2]].values)

    snp2chr = dict(pd.read_csv(data_path + 'misc_data/ukb_white_with_cm.bim', sep = '\t', header = None)[[1, 0]].values)

    idx2clump = {'Clump ' + str(idx): i for idx, i in enumerate(clumps)}

    idx2clump_chromosome = {'Clump ' + str(idx): snp2chr[i[0]] for idx, i in enumerate(clumps)}

    idx2clump_cm = {'Clump ' + str(idx): snp2cm[i[0]] for idx, i in enumerate(clumps)}

    idx2clump_cm_min = {'Clump ' + str(idx): min([snp2cm[x] for x in i]) for idx, i in enumerate(clumps)}

    idx2clump_cm_max = {'Clump ' + str(idx): max([snp2cm[x] for x in i]) for idx, i in enumerate(clumps)}

    df_clumps = pd.DataFrame([idx2clump_chromosome, idx2clump_cm, idx2clump, idx2clump_cm_min, idx2clump_cm_max]).T.reset_index().rename(columns = {'index': 'Clump idx', 0: 'Chromosome', 1: 'cM',  2: 'Clump rsids', 3: 'cM_min',4: 'cM_max'})

    all_mega_clump_across_chr = []

    for chrom in df_clumps.Chromosome.unique():

        df_clump_chr = df_clumps[df_clumps.Chromosome == chrom]

        all_mega_clump = []

        cur_mega_clump = []

        base_cM = 0

        for i,cM_hit,cM_min,cM_max in df_clump_chr.sort_values('cM')[['Clump idx', 'cM', 'cM_min', 'cM_max']].values:

            if (cM_min - base_cM) < 0.1:

                cur_mega_clump.append(i)

                base_cM = cM_max

            else:

                ### this clump is >0.1 cM farther away from the previous clump

                all_mega_clump.append(cur_mega_clump)

                base_cM = cM_max

                cur_mega_clump = [i]

        all_mega_clump.append(cur_mega_clump)

        if len(all_mega_clump[0]) == 0:

            all_mega_clump_across_chr += all_mega_clump[1:]

        else:

            all_mega_clump_across_chr += all_mega_clump

    idx2mega_clump = {'Mega-Clump '+str(idx): i for idx, i in enumerate(all_mega_clump_across_chr)}

    def flatten(l):

        return [item for sublist in l for item in sublist]

    idx2mega_clump_rsid = {'Mega-Clump '+str(idx): flatten([idx2clump[j] for j in i]) for idx, i in enumerate(all_mega_clump_across_chr)}

    idx2mega_clump_chrom = {'Mega-Clump '+str(idx): idx2clump_chromosome[i[0]] for idx, i in enumerate(all_mega_clump_across_chr)}

    return idx2mega_clump, idx2mega_clump_rsid, idx2mega_clump_chrom

def get_mega_clump_query(data_path, clumps, snp_hits, no_hla = False, snp2ld_snps = None):

    snp2ld_snps_with_hla = load_dict(data_path + 'ld_score/ukb_white_ld_10MB.pkl')

    snp2ld_snps_no_hla = load_dict(data_path + 'ld_score/ukb_white_ld_10MB_no_hla.pkl')

    if not snp2ld_snps:

        if no_hla:

            snp2ld_snps = snp2ld_snps_no_hla

        else:

            snp2ld_snps = snp2ld_snps_with_hla

    clumps_pred = []

    snps_in_clumps_pred = []

    K = max(len(clumps) * 3, 100)

    for snp in snp_hits:

        ## top ranked snps

        if len(clumps_pred) >= K:

            ## just going to get the top K clumps where K is set to be very large number -> we don't generate all clumps since as K goes extremely large, they are never prioritized and evaluated.

            break

        else:

            if snp in snps_in_clumps_pred:

                ## already in previous found clumps, move forward

                pass

            else:

                if snp in snp2ld_snps:

                    # this snp has ld tagged snps

                    clumps_pred.append([snp] + snp2ld_snps[snp])

                    snps_in_clumps_pred += snp2ld_snps[snp]

                    snps_in_clumps_pred += [snp]

                else:

                    # this snp does not have ld tagged snps, at least in UKB

                    clumps_pred.append([snp])

                    snps_in_clumps_pred += [snp]

    idx2mega_clump_pred, idx2mega_clump_rsid_pred, idx2mega_clump_chrom_pred = get_meta_clumps(clumps_pred, data_path)

    return idx2mega_clump_pred, idx2mega_clump_rsid_pred, idx2mega_clump_chrom_pred

def get_curve(mega_clump_pred, mega_clump_gold):

    recall_k = {}

    precision_k = {}

    found_clump_idx = []

    clump_idx_record = {}

    pred_clump_has_hit_count = 0

    for k, query_clump in enumerate(mega_clump_pred):

        ## go through the predicted top ranked clumps one by one

        k += 1

        does_this_clump_overlap_with_any_true_clumps = False

        ## this is used to calculate precision, to see if this clump overlaps with any of the gold clumps

        for clump_idx, clump in enumerate(mega_clump_gold):

            ## overlaps with this gold clump

            if len(np.intersect1d(query_clump, clump)) > 0:

                if clump_idx not in found_clump_idx:

                    ## if the clump is never found before, flag it

                    found_clump_idx.append(clump_idx)

                does_this_clump_overlap_with_any_true_clumps = True

        clump_idx_record[k] = copy(found_clump_idx)

        if does_this_clump_overlap_with_any_true_clumps:

            pred_clump_has_hit_count += 1

        recall_k[k] = len(found_clump_idx)/len(mega_clump_gold)

        precision_k[k] = pred_clump_has_hit_count/k

    #sns.scatterplot([recall_k[k+1] for k in range(len(mega_clump_pred))], [precision_k[k+1] for k in range(len(mega_clump_pred))], s = 1)

    return recall_k, precision_k, clump_idx_record

def get_prec_recall(pred_hits, gold_hits):

    recall = len(np.intersect1d(pred_hits, gold_hits))/len(gold_hits)

    if len(pred_hits) != 0:

        precision = len(np.intersect1d(pred_hits, gold_hits))/len(pred_hits)

    else:

        precision = 0

    return {'recall': recall,

           'precision': precision}

def find_nearest(array, value):

    array = np.asarray(array)

    idx = (np.abs(array - value)).argmin()

    return array[idx]

def get_cluster_from_gwas(df, cluster_distance_threshold = 500000, \

                          threshold_extend = False, cluster_compare_threshold = None, \

                         verbose = True):

    cluster_chr_pos = {}

    cluster_chr_rs = {}

    for chr_num in df['#CHROM'].unique():

        df_hits_chr = df[df['#CHROM'] == chr_num]

        df_hits_chr = df_hits_chr.sort_values('POS')

        pos = df_hits_chr.POS.values

        rs = df_hits_chr.ID.values

        cluster_set = []

        cluster_set_rs = []

        cur_pos = pos[0]

        cur_rs = rs[0]

        cur_set = [cur_pos]

        cur_set_rs = [rs[0]]

        for idx, next_pos in enumerate(pos[1:]):

            if next_pos - cur_pos < cluster_distance_threshold:

                cur_set.append(next_pos)

                cur_set_rs.append(rs[idx + 1])

                if threshold_extend:

                    cur_pos = next_pos

            else:

                cluster_set.append(cur_set)

                cluster_set_rs.append(cur_set_rs)

                cur_pos = next_pos

                cur_set = [cur_pos]

                cur_set_rs = [rs[idx + 1]]

        cluster_set.append(cur_set)

        cluster_set_rs.append(cur_set_rs)

        cluster_chr_pos[chr_num] = cluster_set

        cluster_chr_rs[chr_num] = cluster_set_rs

    cluster_chr_pos_flatten = {}

    cluster_chr_cluster_idx_flatten = {}

    cluster_chr_cluster_pos2idx_flatten = {}

    for chr_num, cluster_list in cluster_chr_pos.items():

        pos_flatten = []

        idx_flatten = []

        for idx, cluster in enumerate(cluster_list):

            pos_flatten = pos_flatten + cluster

            idx_flatten = idx_flatten + [idx] * len(cluster)

        cluster_chr_pos_flatten[chr_num] = pos_flatten

        cluster_chr_cluster_idx_flatten[chr_num] = idx_flatten

        cluster_chr_cluster_pos2idx_flatten[chr_num] = dict(zip(pos_flatten, idx_flatten))

    if verbose:

        print('Number of clusters: ' + str(sum([len(j) for j in cluster_chr_pos.values()])))

    cluster_chr_range = {}

    for i,j in cluster_chr_pos.items():

        cluster_chr_range[i] = [(min(x) - cluster_compare_threshold, max(x) + cluster_compare_threshold) for x in j]

    return cluster_chr_pos, cluster_chr_rs, cluster_chr_pos_flatten, \

            cluster_chr_cluster_idx_flatten, cluster_chr_cluster_pos2idx_flatten, cluster_chr_range

def get_cluster_hits_from_pred(pred_hits, threshold, lr_uni, cluster_chr_pos_flatten, cluster_chr_cluster_pos2idx_flatten):

    df_hits = lr_uni[lr_uni.ID.isin(pred_hits)]

    df_hits['closest_cluster'] = df_hits.apply(lambda x: find_nearest(cluster_chr_pos_flatten[x['#CHROM']], x.POS), axis = 1)

    df_hits['distance2cluster'] = df_hits.apply(lambda x: abs(x.closest_cluster - x.POS), axis = 1)

    df_hits['include_as_cluster'] = df_hits.apply(lambda x: x.distance2cluster < threshold, axis = 1)

    df_hits['cluster_id'] = df_hits.apply(lambda x: str(x['#CHROM']) + '_' + str(cluster_chr_cluster_pos2idx_flatten[x['#CHROM']][x['closest_cluster']]), axis = 1)

    cluster2count = dict(df_hits[df_hits.include_as_cluster].cluster_id.value_counts())

    num_non_hits = len(df_hits[~df_hits.include_as_cluster])

    novel_rs_id = df_hits[~df_hits.include_as_cluster].ID.values

    print('Number of predicted hits: ' + str(len(pred_hits)))

    print('Number of predicted hits not in the existing clusters: ' + str(len(novel_rs_id)))

    print('Number of cluster hits: ' + str(len(cluster2count)))

    return cluster2count, num_non_hits, df_hits, novel_rs_id

def plot_cluster_range(chr_num, gnn_cluster_chr_range, cluster_chr_range, \

                       gold_cluster_chr_range, findor_cluster_chr_range, x_start = None, x_end = None, \

                       base_gwas_name = 'FastGWA', gold_ref_name = 'GWAS Catalog'):

    fig = plt.figure(figsize=(14, 3)) # Set the figure size

    ax = fig.add_subplot(111)

    if chr_num not in cluster_chr_range:

        cluster_chr_range[chr_num] = {}

    if chr_num not in gnn_cluster_chr_range:

        gnn_cluster_chr_range[chr_num] = {}

    if chr_num not in gold_cluster_chr_range:

        gold_cluster_chr_range[chr_num] = {}

    if chr_num not in findor_cluster_chr_range:

        findor_cluster_chr_range[chr_num] = {}

    for i in findor_cluster_chr_range[chr_num]:

        plt.plot(i, ['FINDOR', 'FINDOR'], '*-')  

    for i in gnn_cluster_chr_range[chr_num]:

        plt.plot(i, ['GNN', 'GNN'], 's-')

    for i in cluster_chr_range[chr_num]:

        plt.plot(i, [base_gwas_name, base_gwas_name], '^-')

    for i in gold_cluster_chr_range[chr_num]:

        plt.plot(i, [gold_ref_name, gold_ref_name], 'o-')  

    plt.xlabel('Position Index at Chromosome ' + str(chr_num))

    if x_start is not None:

        ax.set_xlim([x_start,x_end])

    plt.show()

def get_pr_curve(cluster_distance_threshold, gold_label_gwas_hits, method_hit_gwas, low_data_gwas_hits, \

                 cluster_compare_threshold = None, method_name = 'gnn'):

    if cluster_compare_threshold is None:

        cluster_compare_threshold = int(cluster_distance_threshold/2)

    gold_cluster_chr_pos, gold_cluster_chr_rs, \

    gold_cluster_chr_pos_flatten, gold_cluster_chr_cluster_idx_flatten, \

    gold_cluster_chr_cluster_pos2idx_flatten, gold_cluster_chr_range = get_cluster_from_gwas(gold_label_gwas_hits, \

                                                                     cluster_distance_threshold, \

                                                                    threshold_extend = threshold_extend, \

                                                                    cluster_compare_threshold = cluster_compare_threshold, \

                                                                    verbose = False)

    cluster_chr_pos, cluster_chr_rs, \

    cluster_chr_pos_flatten, cluster_chr_cluster_idx_flatten, \

    cluster_chr_cluster_pos2idx_flatten, cluster_chr_range = get_cluster_from_gwas(low_data_gwas_hits, \

                                                                cluster_distance_threshold, \

                                                                threshold_extend = threshold_extend, \

                                                                cluster_compare_threshold = cluster_compare_threshold, \

                                                                verbose = False)

    gnn_cluster_chr_pos, gnn_cluster_chr_rs, \

    gnn_cluster_chr_pos_flatten, gnn_cluster_chr_cluster_idx_flatten, \

    gnn_cluster_chr_cluster_pos2idx_flatten, gnn_cluster_chr_range = get_cluster_from_gwas(method_hit_gwas, \

                                                                    cluster_distance_threshold, \

                                                                    threshold_extend = threshold_extend, \

                                                                    cluster_compare_threshold = cluster_compare_threshold, \

                                                                    verbose = False)        

    total = sum([len(j) for i,j in gold_cluster_chr_range.items()])

    #plink_set_overlap = sum([len(j) for j in find_overlap_clusters(cluster_chr_range, gold_cluster_chr_range).values()])

    plink_set_total = sum([len(j) for i,j in cluster_chr_range.items()])

    plink_set_overlap_ref = 0

    plink_set_overlap_query = 0

    for j in find_overlap_clusters(cluster_chr_range, gold_cluster_chr_range).values():

        plink_set_overlap_ref += len(np.unique([set(i[1]) for i in j]))

        plink_set_overlap_query += len(np.unique([set(i[0]) for i in j]))

    #gnn_set_overlap = sum([len(j) for j in find_overlap_clusters(gnn_cluster_chr_range, gold_cluster_chr_range).values()])

    gnn_set_total = sum([len(j) for i,j in gnn_cluster_chr_range.items()])

    gnn_set_overlap_ref = 0

    gnn_set_overlap_query = 0

    for j in find_overlap_clusters(gnn_cluster_chr_range, gold_cluster_chr_range).values():

        gnn_set_overlap_ref += len(np.unique([set(i[1]) for i in j]))

        gnn_set_overlap_query += len(np.unique([set(i[0]) for i in j]))

    '''

    low_data_gold_hits = low_data_gwas[low_data_gwas.ID.isin(gold_label_gwas_hits.ID.values)]

    low_data_gold_hits['cluster_id'] = low_data_gold_hits.apply(lambda x: str(x['#CHROM']) + '_' + \

                                                            str(gold_cluster_chr_cluster_pos2idx_flatten[x['#CHROM']][x.POS]), axis = 1)

    cluster2min_p = dict(low_data_gold_hits.groupby('cluster_id').P.min())

    flat_clusters = [i for i,j in cluster2min_p.items() if j > 1e-3]

    gold_label_gwas_hits['closest_cluster'] = gold_label_gwas_hits.apply(lambda x: find_nearest(gold_cluster_chr_pos_flatten[x['#CHROM']], x.POS), axis = 1)

    gold_label_gwas_hits['distance2cluster'] = gold_label_gwas_hits.apply(lambda x: abs(x.closest_cluster - x.POS), axis = 1)

    gold_label_gwas_hits['cluster_id'] = gold_label_gwas_hits.apply(lambda x: str(x['#CHROM']) + '_' + str(gold_cluster_chr_cluster_pos2idx_flatten[x['#CHROM']][x['closest_cluster']]), axis = 1)

    pos_pred = np.unique(low_data_gwas_hits.ID.values.tolist() + pred_hits.tolist())

    flat_cluster_range = {}

    for i in flat_clusters:

        chr_num = int(i.split('_')[0])

        cluster_idx = int(i.split('_')[1])

        if chr_num in flat_cluster_range:

            flat_cluster_range[chr_num].append(gold_cluster_chr_range[chr_num][cluster_idx])

        else:

            flat_cluster_range[chr_num] = [gold_cluster_chr_range[chr_num][cluster_idx]]

    flat_cluster_recalled = sum([len(j) for j in find_overlap_clusters(gnn_cluster_chr_range, flat_cluster_range).values()])

    flat_cluster_recalled_plink = sum([len(j) for j in find_overlap_clusters(cluster_chr_range, flat_cluster_range).values()])

    '''

    if gnn_set_total == 0:

        gnn_set_precision = -1

    else:

        gnn_set_precision = gnn_set_overlap_query/gnn_set_total

    if plink_set_total == 0:

        plink_precision = -1

    else:

        plink_precision = plink_set_overlap_query/plink_set_total

    return {'plink_precision':plink_precision, 

            'plink_recall': plink_set_overlap_ref/total,

            method_name + '_precision': gnn_set_precision,

            method_name + '_recall': gnn_set_overlap_ref/total,

            'plink_set_overlap_ref': plink_set_overlap_ref,

            'plink_set_overlap_query': plink_set_overlap_query,

            'plink_set_total': plink_set_total,

            method_name + '_set_overlap_ref': gnn_set_overlap_ref,

            method_name + '_set_overlap_query': gnn_set_overlap_query,

            method_name + '_set_total': gnn_set_total,

            'total_set': total

            #'gnn_flat_cluster_recall': flat_cluster_recalled/len(flat_clusters),

            #'plink_flat_cluster_recall': flat_cluster_recalled_plink/len(flat_clusters)

           }

from tqdm import tqdm

def find_overlap_clusters(query_cluster2range, gold_cluster2range):

    set_found_cluster_all = {}

    for chr_num, eval_cluster in query_cluster2range.items():

        if chr_num in gold_cluster2range:

            gold_cluster = gold_cluster2range[chr_num]

            set_found_cluster = []

            for a in eval_cluster:

                for b in gold_cluster:

                    if (a[0] <= b[1]) and (b[0] <= a[1]):

                        set_found_cluster.append((a, b))

                        break

            set_found_cluster_all[chr_num] = set_found_cluster 

    return set_found_cluster_all

def find_non_overlap_clusters(query_cluster2range, gold_cluster2range):

    set_not_found_cluster_all = {}

    for chr_num, eval_cluster in query_cluster2range.items():

        gold_cluster = gold_cluster2range[chr_num]

        set_not_found_cluster = []

        for a in eval_cluster:

            set_found_cluster = []

            for b in gold_cluster:

                if (a[0] <= b[1]) and (b[0] <= a[1]):

                    set_found_cluster.append((a, b))

                    break

            if len(set_found_cluster) == 0:

                set_not_found_cluster.append(a)

        set_not_found_cluster_all[chr_num] = set_not_found_cluster 

    return set_not_found_cluster_all

### eval support functions

def quantileNormalize(df_input):

    df = df_input.copy()

    #compute rank

    dic = {}

    for col in df:

        dic.update({col : sorted(df[col])})

    sorted_df = pd.DataFrame(dic)

    rank = sorted_df.mean(axis = 1).tolist()

    #sort

    for col in df:

        t = np.searchsorted(np.sort(df[col]), df[col])

        df[col] = [rank[i] for i in t]

    return df

def get_cluster_count(method_hit_gwas, cluster_distance_threshold, cluster_compare_threshold, threshold_extend, gold_cluster_chr_range):

    gnn_cluster_chr_pos, gnn_cluster_chr_rs, \

    gnn_cluster_chr_pos_flatten, gnn_cluster_chr_cluster_idx_flatten, \

    gnn_cluster_chr_cluster_pos2idx_flatten, gnn_cluster_chr_range = get_cluster_from_gwas(method_hit_gwas, \

                                                                    cluster_distance_threshold, \

                                                                    threshold_extend = threshold_extend, \

                                                                    cluster_compare_threshold = cluster_compare_threshold, \

                                                                    verbose = False)        

    total = sum([len(j) for i,j in gold_cluster_chr_range.items()])

    gnn_set_total = sum([len(j) for i,j in gnn_cluster_chr_range.items()])

    gnn_set_overlap_ref = 0

    gnn_set_overlap_query = 0

    for j in find_overlap_clusters(gnn_cluster_chr_range, gold_cluster_chr_range).values():

        gnn_set_overlap_ref += len(np.unique([set(i[1]) for i in j]))

        gnn_set_overlap_query += len(np.unique([set(i[0]) for i in j]))

    return {'set_overlap_ref': gnn_set_overlap_ref,

            'set_overlap_query': gnn_set_overlap_query,

            'set_total': gnn_set_total,

            'total_set': total

           }

## search every 100 until it is larger than k, then search every 10, then search every 1

def get_top_k_clusters(query_rank, top_hits_k_range, cluster_distance_threshold, cluster_compare_threshold, threshold_extend, gold_cluster_chr_range):

    snp_k = 0

    k_to_cluster = {}

    k_to_closest_x = {}

    for k in top_hits_k_range:

        while True:

            out = get_cluster_count(query_rank[:snp_k], cluster_distance_threshold, 

                          cluster_compare_threshold, threshold_extend, gold_cluster_chr_range)

            if out['set_total'] < k:

                snp_k += 100

            else:

                snp_k -= 100

                while True:

                    out = get_cluster_count(query_rank[:snp_k], cluster_distance_threshold, 

                          cluster_compare_threshold, threshold_extend, gold_cluster_chr_range)

                    if out['set_total'] < k:

                        snp_k += 10

                    else:

                        closest_x = snp_k

                        closest_distance = abs(out['set_total'] - k)

                        for x in range(snp_k - 10, snp_k):

                            out = get_cluster_count(query_rank[:x], cluster_distance_threshold, 

                                  cluster_compare_threshold, threshold_extend, gold_cluster_chr_range)

                            if abs(out['set_total'] - k) <= closest_distance:

                                closest_x = x

                                closest_distance = abs(out['set_total'] - k)

                        break

                break

        k_to_cluster[k] = get_cluster_count(query_rank[:closest_x], cluster_distance_threshold, 

                      cluster_compare_threshold, threshold_extend, gold_cluster_chr_range)

        k_to_closest_x[k] = closest_x

    return k_to_cluster, k_to_closest_x

def storey_pi_estimator(gwas_data, bin_index):

    """

    Estimate pi0/pi1 using Storey and Tibshirani (PNAS 2003) estimator.

    Argss

    =====

    bin_index: array of indices for a particular bin

    """

    pvalue = gwas_data.loc[bin_index,'P'] # extract pvalues from specific bin based index

    #assert(pvalue.min() >= 0 and pvalue.max() <= 1), "Error: p-values should be between 0 and 1"

    total_tests = float(len(pvalue))

    pi0 = []

    lam = np.arange(0.05, 0.95, 0.05)

    counts = np.array([(pvalue > i).sum() for i in np.arange(0.05, 0.95, 0.05)])

    for l in range(len(lam)):

        pi0.append(counts[l] / (total_tests * (1 - lam[l])))

    # fit  cubic spline

    if not np.all(np.isfinite(pi0)):

        print("Not all pi0 is finite!!! filtering to finite indices...")

        finite_indices = np.isfinite(pi0)

        lam = lam[finite_indices]

        pi0 = pi0[finite_indices]

    cubic_spline = interpolate.CubicSpline(lam, pi0)

    pi0_est = cubic_spline(lam[-1])

    if(pi0_est >1): #take care of out of bounds estimate

        pi0_est = 1

    return pi0_est

def storey_ribshirani_integrate(gwas_data, column = 'pred', num_bins = 100):

    num_bins = float(num_bins)

    quantiles = np.arange(0, 1 + 1 / (num_bins+1), 1 / num_bins)

    predicted_tagged_variance_quantiles = gwas_data[column].quantile(quantiles)

    #expand top quantiles to ensure everything is within range

    predicted_tagged_variance_quantiles[0] = predicted_tagged_variance_quantiles[0]-1

    predicted_tagged_variance_quantiles[1] = predicted_tagged_variance_quantiles[1]+1

    predicted_tagged_variance_quantiles = predicted_tagged_variance_quantiles.drop_duplicates()

    num_bins = len(predicted_tagged_variance_quantiles)-1

    bins = pd.cut(gwas_data[column], predicted_tagged_variance_quantiles, labels=np.arange(num_bins)) #create the lables

    gwas_data['bin_number'] = bins

    gwas_data['pi0'] = None

    if (gwas_data['P'].min() < 0) or (gwas_data['P'].max() > 1):

        print("detected p-values < 0 or > 1, please double check. we clipped it to 0-1 for now...")

        gwas_data['P'] = gwas_data['P'].clip(lower=0, upper=1)

    #print("Estimating pi0 within each bin")

    for i in range(num_bins):

        bin_index = gwas_data['bin_number']== i # determine index of snps in bin number i

        if len(gwas_data[bin_index])>0:

            pi0 = storey_pi_estimator(gwas_data, bin_index)

            ## preventing exploding weights

            if pi0 < 1e-5:

                pi0 = 1e-5

            if pi0 > 1-1e-5:

                pi0 = 1-1e-5

            gwas_data.loc[bin_index, 'pi0'] = pi0

    if any(gwas_data['pi0'] == 1): # if a bin is estimated to be all null, give the smallest non-null weight

        one_index = gwas_data['pi0'] == 1

        largest_pi0 = gwas_data.loc[~one_index]['pi0'].max()

        gwas_data.loc[one_index,'pi0'] = largest_pi0

    if any(gwas_data['pi0'] == 0): # if a bin is estimated to be all alternative, give the largest non-null weight

        one_index = gwas_data['pi0'] == 0

        largest_pi0 = gwas_data.loc[~one_index]['pi0'].min()

        gwas_data.loc[one_index,'pi0'] = largest_pi0

    #print("Re-weighting SNPs")

    weights = (1-gwas_data['pi0'])/(gwas_data['pi0'])

    ## avoiding exploding p-values

    #weights = np.maximum(1, weights.values)

    mean_weight = weights.mean()

    weights = weights/mean_weight #normalize weights to have mean 1

    ## avoiding exploding p-values

    #weights = np.maximum(1, weights.values)

    gwas_data['weights'] = weights

    gwas_data['P_weighted'] = gwas_data['P']/weights #reweight SNPs

    index = gwas_data['P_weighted'] > 1

    #gwas_data.loc[index, 'P_weighted'] = 1

    gwas_data.loc[index, 'P_weighted'] = gwas_data['P'][index] ## using original p-value when above 1

    gwas_data.loc[gwas_data['P_weighted'].isnull(), 'P_weighted'] = 1    

    return gwas_data['P_weighted'].values