"""FDR correction for U-shaped p-value distributions as proposed by
Pounds et al., Bioinformatics 2006.

Copyright (c) 2015 Johannes Koester, Xiaole Liu lab
This code is free software; you can redistribute it and/or modify it
under the terms of the BSD License (see the file COPYING included with
the distribution).
@status:  experimental
@version: $Revision$
@author:  Johannes Koester
@contact: koester@jimmy.harvard.edu
"""



from itertools import groupby
from bisect import bisect


def est_null_prop_continous(pvalues):
    """
    Equation 6 and 7 in Pounds et al., Bioinformatics 2006.
    """
    a = [2 * min(p, 1 - p) for p in pvalues]
    null_prop = min(1, 2 * sum(a) / len(a))
    return null_prop


def est_significant_prop(sorted_pvalues):
    """
    Estimate expected propotion of p-values smaller or equal to each possible
    alpha empirically from given p-values. Equation (5) in
    Pounds et al., Bioinformatics 2006.
    """
    sig_prop = []
    total = 0
    for p, equal in groupby(sorted_pvalues):
        k = sum(1 for p in equal)
        total += k
        sig_prop.extend([total / len(sorted_pvalues)] * k)
    return sig_prop


def fdr_pounds_continuous(pvalues):
    """
    Calculate FDR-adjusted p-values with the method described in
    Pounds et al., Bioinformatics 2006.
    This function assumes that p-values approximate a continuous distribution.
    """
    idx = sorted(list(range(len(pvalues))), key=pvalues.__getitem__)
    sorted_pvalues = [pvalues[i] for i in idx]

    null_prop = est_null_prop_continous(pvalues)
    sig_prop = est_significant_prop(sorted_pvalues)

    # get the first p-value >0.5 by binary search
    middle = bisect(sorted_pvalues, 0.5)

    # adjust p-values <= 0.5
    # the p-value is adjusted by the estimated proportion of true null hypotheses
    lfdr = [min(1.0, null_prop * p / s)
            for p, s in zip(sorted_pvalues[:middle], sig_prop[:middle])]
    # adjust p-values > 0.5
    # here, the assumption is that all p-values are generated by true null hypotheses
    lfdr.extend(min(1.0, (null_prop / 2 + s - sig_prop[middle - 1]) / s)
                for s in sig_prop[middle:])

    # initialize FDR array
    fdr = [1] * len(pvalues)
    last = 1
    for i, t in zip(reversed(idx), reversed(lfdr)):
        fdr[i] = min(t, last)
        last = fdr[i]

    return fdr


# calculate the FDR of p values using false discovery rate
#def pFDR(x):
#    index = [i[0] for i in sorted(enumerate(x), key=lambda v:v[1],reverse=True)]
#    rindex = [i[0] for i in sorted(enumerate(index), key=lambda v:v[1])]
#    qval = sum([1.0/i for i in xrange(1,len(x)+1)])
#    l = [qval*len(x)/i*x[j] for i,j in zip(reversed(xrange(1,len(x)+1)),index)]
#    l = [l[k] if l[k] < 1.0 else 1.0 for k in rindex]
#    return l

def pFDR(pvalues, method='fdr'):
  """
  calculate the FDR of p values using false discovery rate
  check http://stackoverflow.com/questions/7450957/how-to-implement-rs-p-adjust-in-python
  """
  n=len(pvalues)
  new_pvalues=[0]*n

  if method=='holm':
    values = [ (pvalue, i) for i, pvalue in enumerate(pvalues) ]
    values.sort()
    for rank, vals in enumerate(values):
      pvalue, i = vals
      new_pvalues[i] = (n-rank) * pvalue
      if new_pvalues[i] > 1.0:
        new_pvalues[i]=1.0
  elif method=='fdr':
    values = [ (pvalue, i) for i, pvalue in enumerate(pvalues) ]
    values.sort()
    values.reverse()
    new_values = []
    for i, vals in enumerate(values):
      rank = n - i
      pvalue, index = vals
      new_values.append((n/rank) * pvalue)
    for i in range(0, int(n)-1):
      if new_values[i] < new_values[i+1]:
        new_values[i+1] = new_values[i]
    for i, vals in enumerate(values):
      pvalue, index = vals
      new_pvalues[index] = new_values[i]
  elif method=='pounds':
    new_pvalues=fdr_pounds_continuous(pvalues)
  else:
    logging.error('Unknown FDR mthod '+method)
    return pvalues
  return new_pvalues



if __name__ == "__main__":
    # take gene_summary from stdin and plot some measures to ensure that the method makes sense
    import sys
    import pandas as pd
    import matplotlib.pyplot as plt
    data = pd.read_table(sys.stdin)
    for sel in "pos neg".split():
        d = data[["p." + sel, "fdr." + sel]].copy()
        d.columns = ["p", "fdr"]
        d.sort("p", inplace=True)
        pvalues = list(d["p"].values)
        d["fdr.pounds"] = fdr_continuous(pvalues)

        plt.figure(figsize=(15, 15))
        plt.suptitle("selection: " + sel)
        plt.subplot(221)
        plt.plot(d["p"], d["fdr.pounds"], "r-", label="Pounds")
        plt.plot(d["p"], d["fdr"], "k-", label="BH")
        plt.xlabel("p-values")
        plt.ylabel("FDR")
        plt.legend()

        plt.subplot(222)
        plt.plot(d["fdr"], d["fdr.pounds"], "k-")
        plt.plot([0, 1], [0, 1], "--", color="grey")
        plt.xlabel("FDR (BH)")
        plt.ylabel("FDR (Pounds)")

        plt.subplot(223)
        d["p"].hist(normed=True,
                    cumulative=True,
                    histtype="step",
                    color="blue",
                    bins=100,
                    label="p-value")
        d["fdr"].hist(normed=True,
                      cumulative=True,
                      histtype="step",
                      bins=100,
                      color="black",
                      label="FDR (BH)")
        d["fdr.pounds"].hist(normed=True,
                             cumulative=True,
                             histtype="step",
                             bins=100,
                             color="red",
                             label="FDR (Pounds)")
        plt.ylabel("cumulative density")
        plt.legend(loc="upper left")
        plt.grid(False)

        plt.subplot(224)
        d["p"].hist(normed=True, color="black", bins=100)
        null_prop = est_null_prop_continous(pvalues)
        plt.plot([0, 1], [null_prop] * 2, "r--",
                 label="estimated true null hypotheses (Pounds)")
        plt.legend()
        plt.xlabel("p-value")
        plt.ylabel("density")
        plt.grid(False)
    plt.show()
