try:
    import rpy2.robjects.numpy2ri
    from rpy2.robjects import r as r
    from rpy2 import rinterface
except ImportError:
    pass

import numpy as np
import scipy as sp
import sys, pickle, pdb
import scipy.stats as st
import scipy.interpolate
import logging as LG

def R_qvalues(pv):
    rinterface.set_writeconsole(None)
    r.library("qvalue")
    bench = r("qvalue")
    
    results = bench(pv)
    pi0 = results[1]
    qv = np.array(results[2]) 

    return qv, pi0 


def qvalues1(PV,m=None,pi=1.0):
    """estimate q vlaues from a list of Pvalues
    this algorihm is taken from Storey, significance testing for genomic ...
    m: number of tests, (if not len(PV)), pi: fraction of expected true null (1.0 is a conservative estimate)
    """
          
    S = PV.shape
    PV = PV.flatten()
    if m is None:
        m = len(PV) * 1.0
    else:
        m*=1.0
    lPV = len(PV)
    
    #1. sort pvalues
    PV = PV.squeeze()
    IPV = PV.argsort()
    PV  = PV[IPV]

    #2. estimate lambda
    if pi is None:
        lrange = sp.linspace(0.05,0.95,max(lPV/100.0,10))
        pil    = sp.double((PV[:,sp.newaxis]>lrange).sum(axis=0))/lPV
        pilr   = pil/(1.0-lrange)
        #ok, I think for SNPs this is pretty useless, pi is close to 1!
        pi =1.0
        #if there is something useful in there use the something close to 1
        if pilr[-1]<1.0:
            pi = pilr[-1]
            
    #3. initialise q values
    QV_ = pi * m/lPV* PV
    QV_[-1] = min(QV_[-1],1.0)
    #4. update estimate
    for i in xrange(lPV-2,-1,-1):
        QV_[i] = min(pi*m*PV[i]/(i+1.0),QV_[i+1])
    #5. invert sorting
    QV = sp.zeros_like(PV)
    QV[IPV] = QV_

    QV = QV.reshape(S)
    return QV
    
def qvalues(pv, m = None, return_pi0 = False, lowmem = False, pi0 = None, fix_lambda = None):

    original_shape = pv.shape

    assert(pv.min() >= 0 and pv.max() <= 1)

    pv = pv.ravel() # flattens the array in place, more efficient than flatten() 

    if m == None:
	m = float(len(pv))
    else:
	# the user has supplied an m, let's use it
	m *= 1.0

    # if the number of hypotheses is small, just set pi0 to 1
    if len(pv) < 100:
	pi0 = 1.0
    elif pi0 != None:
	pi0 = pi0
    else:
	# evaluate pi0 for different lambdas
	pi0 = []
	lam = sp.arange(0, 0.90, 0.01)
	counts = sp.array([(pv > i).sum() for i in sp.arange(0, 0.9, 0.01)])
	
	if fix_lambda != None:
	    interv_count = (pv > fix_lambda - 0.01).sum()
	    uniform_sim = sp.array([(pv > fix_lambda-0.01).sum()*(i+1) for i in sp.arange(0, len(sp.arange(0, 0.90, 0.01)))][::-1])
	    counts += uniform_sim
	    
	for l in range(len(lam)):
	    pi0.append(counts[l]/(m*(1-lam[l])))

	pi0 = sp.array(pi0)

	# fit natural cubic spline
	tck = sp.interpolate.splrep(lam, pi0, k = 3)
	pi0 = sp.interpolate.splev(lam[-1], tck)
	if pi0 > 1:
	    LG.warning("got pi0 > 1 (%.3f) while estimating qvalues, setting it to 1" % pi0)
	    pi0 = 1.0

	assert(pi0 >= 0 and pi0 <= 1), "%f" % pi0


    if lowmem:
	# low memory version, only uses 1 pv and 1 qv matrices
	qv = sp.zeros((len(pv),))
	last_pv = pv.argmax()
	qv[last_pv] = (pi0*pv[last_pv]*m)/float(m)
	pv[last_pv] = -sp.inf
	prev_qv = last_pv
	for i in xrange(int(len(pv))-2, -1, -1):
	    cur_max = pv.argmax()
	    qv_i = (pi0*m*pv[cur_max]/float(i+1))
	    pv[cur_max] = -sp.inf
	    qv_i1 = prev_qv
	    qv[cur_max] = min(qv_i, qv_i1)
	    prev_qv = qv[cur_max]

    else:
	p_ordered = sp.argsort(pv)    
	pv = pv[p_ordered]
	# estimate qvalues
# 	qv = pi0*m*pv/(sp.arange(len(pv))+1.0)
	
# 	for i in xrange(int(len(qv))-2, 0, -1):
# 	    qv[i] = min([qv[i], qv[i+1]])


	qv = pi0 * m/len(pv) * pv
	qv[-1] = min(qv[-1],1.0)

	for i in xrange(len(pv)-2, -1, -1):
	    qv[i] = min(pi0*m*pv[i]/(i+1.0), qv[i+1])



	# reorder qvalues
	qv_temp = qv.copy()
	qv = sp.zeros_like(qv)
	qv[p_ordered] = qv_temp



    # reshape qvalues
    qv = qv.reshape(original_shape)

    if return_pi0:
	return qv, pi0
    else:
	return qv


def estimate_lambda(pv):
    """estimate lambda form a set of PV"""
    LOD2 = sp.median(st.chi2.isf(pv,1))
    L = (LOD2/0.456)
    return (L)

def LOD2PV(lods):
    PV = (st.chi2.sf(2*lods, 1))
    return PV
        
if __name__ == '__main__':
    qv1, pvalues = pickle.load(open("../experiments_limmi/simulation/results_001_nocorr_0/PANAMA.pickle", "r"))
    # pvalues = pvalues[:20, :100]
    qv = qvalues(pvalues, lowmem = True)
    Rqv, Rpi0 = R_qvalues(pvalues)

    print np.sqrt(sp.power((qv - Rqv), 2).mean())

