## modules directed towards enrichment of node partitions..

##### this pyton code enables enrichment calculation from graph results from previous step

## this is to calculate enrichment scores

from scipy.stats import fisher_exact
import multiprocessing as mp
import random
from statsmodels.sandbox.stats.multicomp import multipletests
from collections import defaultdict, Counter
from ..term_parsers import parse_gaf_file,read_termlist,read_topology_mappings,read_uniprot_GO
import pandas as pd
import numpy as np
import logging
logging.basicConfig(format='%(asctime)s - %(message)s', datefmt='%d-%b-%y %H:%M:%S')
logging.getLogger().setLevel(logging.INFO)

def calculate_pval(term,alternative="two-sided"):

    """
    Parallel kernel for computation of p vals. All partitions are considered with respect to agiven GO term! Counts in a given partition are compared to population.

    """
    
    query_term = term[0]
    query_term_count_population = term[1]

    inside_local = 0
    outside_local = 0
    for x in _partition_entries:
        terms = _map_term_database[x]
        if query_term in terms:
            inside_local+=1
        else:
            outside_local+=1

    if _normalize_by_comsize:
        query_term_count_population = int(query_term_count_population/_number_of_communities)
            
    ## EASE correction.
    if inside_local > 1:
        query_counts = [inside_local-1, query_term_count_population]
    else:
        query_counts = [inside_local, query_term_count_population]
    pop_counts = [outside_local, _number_of_all_annotated-query_term_count_population]
    p_value = fisher_exact([query_counts,pop_counts],alternative=alternative)[1]
    return p_value

def multiple_test_correction(input_dataset):

    """
    Multiple test correction. Given a dataset with corresponding significance levels, perform MTC.
    """
    
    from statsmodels.sandbox.stats.multicomp import multipletests
    pvals = defaultdict(list)
    with open(input_dataset) as ods:
        for line in ods:
            try:
                component, term, pval = line.split()
                pvals[component].append((term,pval))
            except:
                pass

    logging.info("Component_by_size PFAM_term pvalue")
    for key, values in pvals.items():
        tmpP = [float(val[1]) for val in values]
        termN = [val[0] for val in values]
        significant, pvals, sidak, bonf = multipletests(tmpP,method="hs",is_sorted=False,returnsorted=False)

        ## Holm Sidak        
        output = zip(termN,significant,pvals,tmpP)
        for term,significant,pval,tmp in output:
            if (significant == True):
                print (key,term,significant,tmp,pval)

def parallel_enrichment(term):

    """
    A simple kernel for parallel computation of p-values. (assuming independence of experiments)
    """
    
    pval = calculate_pval(_term_database[term],alternative=_alternative)
    return {'observation' : _partition_name,'term' : _term_database[term][0],'pval' : pval}

def compute_enrichment(term_dataset, term_database, topology_map, all_counts, whole_term_list=False,pvalue=0.05,multitest_method="fdr_bh",alternative="two-sided",intra_community=False):

    """
    The main method for computing the enrichment of a subnetwork. This work in parallel and also offers methods for multiple test correction.
    """
    
    if whole_term_list:
        tvals = set.union(*[x for x in topology_map.values()])
        topology_map = {}
        topology_map['1_community'] = tvals
    
    global _partition_name
    global _alternative
    global _partition_entries
    global _term_database
    global _map_term_database
    global _number_of_all_annotated
    global _number_of_communities
    global _normalize_by_comsize
    
    _alternative = alternative
    _number_of_all_annotated = all_counts
    _term_database = {en : x for en, x in enumerate(term_database.items())} ## database of all annotations
    
    _map_term_database = term_dataset ## entry to acc mappings        
    _number_of_communities = len(topology_map)
    _normalize_by_comsize = intra_community

    if intra_community:
        _number_of_all_annotated = int(_number_of_all_annotated/_number_of_communities)
    
    finalFrame = pd.DataFrame()
    
    for k, v in topology_map.items():

        logging.info("Computing enrichment for partition {}".format(k))
        ## reassign for parallel usage
        _partition_name = k
        _partition_entries = v

        ## computational pool instantiation
        ncpu = 2 #mp.cpu_count()
#        pool = mp.Pool(ncpu)
        
        ## compute the results
        n = len(term_database)
        step = ncpu ## number of parallel processes
        jobs = [range(n)[i:i + step] for i in range(0, n, step)] ## generate jobs

        ## result container
        tmpframe = pd.DataFrame(columns=['observation','term','pval'])
        results = [parallel_enrichment(x) for x in range(n)]
        
        # for batch in jobs:
        #     results = pool.map(parallel_enrichment,batch)
        tmpframe = tmpframe.append(results,ignore_index=True)

        ## multitest corrections on partition level
        if multitest_method == "raw":
            tmpframe = tmpframe[tmpframe['pval'] < pvalue]
        else:
            significant, p_adjusted, sidak, bonf = multipletests(tmpframe['pval'],method=multitest_method,is_sorted=False, returnsorted=False, alpha=pvalue)
            tmpframe['corrected_pval'+"_"+multitest_method] = pd.Series(p_adjusted)
            tmpframe['significant'] = pd.Series(significant)            
            tmpframe = tmpframe[tmpframe['significant'] == True]
            
        finalFrame = finalFrame.append(tmpframe,ignore_index=True)
    
    return finalFrame

## specifiy how this is formatted..
def fet_enrichment_generic(term_dataset,term_database,all_counts,topology_map):
    """
    A generic enrichment method useful for arbitrary partition enrichment (CBSSD baseline).

    term_dataset = defaultdict(list) of node:[a1..an] mappings
    term_datset = Counter object of individual annotation occurrences
    all_counts = number of all annotation occurences

    """
    ## 3.) calculate p-vals.
    significant_results = compute_enrichment(term_dataset, term_database, topology_map, all_counts,whole_term_list=False)
    return significant_results

def fet_enrichment_terms(partition_mappings,annotation_mappings,alternative="two-sided",intra_community = False,pvalue=0.1,multitest_method="fdr_bh"):

    """
    This is the most generic enrichment process.
    """
    
    ## 1.) read the database.
    term_dataset, term_database, all_counts =  read_uniprot_GO(annotation_mappings)
    
    ## 2.) partition function dict.
    topology_map = read_topology_mappings(partition_mappings)

    ## 3.) calculate p-vals.
    significant_results = compute_enrichment(term_dataset, term_database, topology_map, all_counts,whole_term_list=False,alternative=alternative,intra_community=intra_community,pvalue=pvalue,multitest_method=multitest_method)

    return significant_results

## write this so it uses vanilla data structures
def fet_enrichment_uniprot(partition_mappings,annotation_mappings):

    """
    This is a pre-designed wrapper for uniprot-like annotation.
    """
    
    ## 1.) read the database.
    term_dataset, term_database, all_counts =  read_uniprot_GO(annotation_mappings)
    
    ## 2.) partition function dict.
    topology_map = read_topology_mappings(partition_mappings)

    ## 3.) calculate p-vals.
    significant_results = compute_enrichment(term_dataset, term_database, topology_map, all_counts,whole_term_list=False)

    return significant_results

if __name__ == "__main__":

    pass
