import networkx as nx
from scipy.stats import hypergeom
from statsmodels.stats.multitest import fdrcorrection

import goenrich.export

def analyze(H, query, **kwargs):
    """ run enrichment analysis for query

    >>> O = goenrich.obo.graph('...')
    >>> background = goenrich.read.goa('...')
    >>> H = goenrich.enrich.set_background(O, background, ...)
    >>> df = goenrich.enrich.analyze(H, query, ...)

    :param H: Ontology graph after backgroud was set
    :param query: array like of ids
    :returns: pandas.DataFrame with results
    """
    options = {
            'node_filter' : lambda node : 'p' in node,
            'show' : 'top20',
            'inplace' : True
    }
    options.update(kwargs)
    if not options['inplace']:
        raise ValueError('goenrich.enrich.analyze needs inplace=False')
    pvalues = calculate_pvalues(H, query, **options)
    multiple_testing_correction(H, pvalues, **options)
    df = goenrich.export.to_frame(H, **options)
    if 'gvfile' in options:
        show = options['show']
        if show.startswith('top'):
            top = int(show.replace('top', ''))
            sig = df.sort('q').head(top).index
        else:
            raise NotImplementedError(show)
        goenrich.export.to_graphviz(H, sig, **options)
    return df
    
def set_background(O, df, entry_id, category_id, background_name='', inplace=False):
    """ Creates a copy of the ontology tree and
    propagates the background set through it

    >>> O = goenrich.obo.graph('...')
    >>> background = goenrich.read.goa('...')
    >>> H = goenrich.enrich.set_background(O, background, ...)
    
    :param O: ontology graph: generated by geonrich.obo.graph
    :param df: background data set: look in goenrich.read for parsers
    :param entry_id: protein or gene identifier column
    :param category_id: GO term column
    :param background_name: specify background annotation name
    :param inplace: force in-place annotation
    :returns H: the annotated ontology graph if inplace was not set
    """
    H = O if inplace else O.copy()
    M = len(df[entry_id].unique()) # total number of objects
    if inplace: # clean background attribute for changing backgrounds
        for n in H:
            node = H.node[n]
            node['background'] = set([])

    grouped = df.groupby(category_id)[entry_id]
    for term,entries in grouped:
        namespace = H.node[term]['namespace']
        root = H.graph['roots'][namespace]
        for path in nx.simple_paths.all_simple_paths(H, term, root):
            for n in path:
                node = H.node[n]
                node.setdefault('background', set()).update(entries)
                if 'M' not in node:
                    node['M'] = M
    H.graph['background'] = background_name
    if not inplace:
        return H

def calculate_pvalues(H, query, min_hit_size=2, min_category_size=3,
        max_category_size=500, max_category_depth=5,
        inplace=False, **kwargs):
    """ calculate pvalues for all categories in the graph
    
    :param H: ontology graph after background was set
    :param query: array_like of identifiers
    :param min_hit_size: minimum intersection size of query and category 
    :param min_category_size: categories smaller than this number are ignored
    :param max_category_size: categories larger than this number are ignored
    :param inplace: perform operations inplace or return copy
    :returns: dictionary of term : pvalue
    """
    G = H if inplace else H.copy()
    query_set = set(query)
    pvalues = {}
    N = len(query_set)
    for i in G:
        node = G.node[i]
        if inplace:
            # reset all query related attributes
            for attr in ['query', 'n', 'N', 'hits', 'x', 'p', 'q', 'significant']:
                if attr in node:
                    del node[attr]
        background = node.get('background', set([]))
        n = len(background)
        node['n'] = n
        hits = query_set.intersection(background)
        x = len(hits)
        depth = node.get('depth', -1) # depth might not be set due to malformed ontology
        if ((depth > max_category_depth)
            or (n < min_category_size)
            or (n > max_category_size)
            or (x < min_hit_size)):
            continue
        else:
            node['query'] = query_set
            node['N'] = N
            node['hits'] = hits
            node['x'] = x
            M, n = node['M'], node['n']
            p = hypergeom.sf(x, M, n, N)

            node['p'] = p
            pvalues[i] = p
    if inplace:
        return pvalues
    else:
        return G, pvalues


def multiple_testing_correction(G, pvalues, alpha=0.05,
        method='benjamini-hochberg', inplace=False, **kwargs):
    """ correct pvalues for multiple testing and add corrected `q` value
    
    :param G: ontology graph for which pvalues were calculated
    :param alpha: significance level default : 0.05
    :param method: multiple testing correction method [bonferroni|benjamini-hochberg]
    """
    H = G if inplace else G.copy()
    H.graph.update({ 'multiple-testing-correction': method, 'alpha' : alpha })
    if method == 'bonferroni':
        n = len(pvalues.values())
        for term,p in pvalues.items():
            node = H.node[term]
            q = p * n
            node['q'] = q
            node['significant'] = q < 0.05
    elif method == 'benjamini-hochberg':
        terms, ps = zip(*pvalues.items())
        rejs, qs = fdrcorrection(ps, alpha)
        for term, q, rej in zip(terms, qs, rejs):
            node = H.node[term]
            node['q'] = q
            node['significant'] = rej
    else:
        raise ValueError(method)
    if not inplace:
        return H
