#!/usr/bin/env python
"""Translations of functions from Release 2.3 of the Cephes Math Library, 
which is (c) Stephen L. Moshier 1984, 1995.
"""
from __future__ import division
from cogent.maths.stats.special import erf, erfc, igamc, igam, betai, log1p, \
    expm1, SQRTH, MACHEP, MAXNUM, PI, ndtri, incbi, igami, fix_rounding_error,\
    ln_binomial
    #ndtri import b/c it should be available via this module

from numpy import sqrt, exp, arctan as atan

__author__ = "Rob Knight"
__copyright__ = "Copyright 2007-2012, The Cogent Project"
__credits__ = ["Rob Knight", "Sandra Smit", "Gavin Huttley", "Daniel McDonald"]
__license__ = "GPL"
__version__ = "1.5.3"
__maintainer__ = "Rob Knight"
__email__ = "rob@spot.colorado.edu"
__status__ = "Production"

incbet = betai  #shouldn't have renamed it...

#Probability integrals: low gives left-hand tail, high gives right-hand tail.
def z_low(x):
    """Returns left-hand tail of z distribution (0 to x). 
    
    x ranges from -infinity to +infinity; result ranges from 0 to 1
    
    See Cephes docs for details."""
    y = x * SQRTH
    z = abs(y) #distribution is symmetric
    if z < SQRTH:
        return 0.5 + 0.5 * erf(y)
    else:
        if y > 0:
            return 1 - 0.5 * erfc(z)
        else:
            return 0.5 * erfc(z)

def z_high(x):
    """Returns right-hand tail of z distribution (0 to x). 
    
    x ranges from -infinity to +infinity; result ranges from 0 to 1
    
    See Cephes docs for details."""
    y = x * SQRTH
    z = abs(y)
    if z < SQRTH:
        return 0.5 - 0.5 * erf(y)
    else:
        if x < 0:
            return 1 - 0.5 * erfc(z)
        else:
            return 0.5 * erfc(z)

def zprob(x):
    """Returns both tails of z distribution (-inf to -x, inf to x)."""
    return 2 * z_high(abs(x))

def chi_low(x, df):
    """Returns left-hand tail of chi-square distribution (0 to x), given df.

    x ranges from 0 to infinity.
    
    df, the degrees of freedom, ranges from 1 to infinity (assume integers).
    Typically, df is (r-1)*(c-1) for a r by c table.
   
    Result ranges from 0 to 1.
    
    See Cephes docs for details.
    """
    x = fix_rounding_error(x)
    if x < 0:
        raise ValueError, "chi_low: x must be >= 0 (got %s)." % x
    if df < 1:
        raise ValueError, "chi_low: df must be >= 1 (got %s)." % df
    return igam(df/2, x/2)
    
def chi_high(x, df):
    """Returns right-hand tail of chi-square distribution (x to infinity).
    
    df, the degrees of freedom, ranges from 1 to infinity (assume integers).
    Typically, df is (r-1)*(c-1) for a r by c table.
    
    Result ranges from 0 to 1.
    
    See Cephes docs for details.
    """
    x = fix_rounding_error(x)
    
    if x < 0:
        raise ValueError, "chi_high: x must be >= 0 (got %s)." % x
    if df < 1:
        raise ValueError, "chi_high: df must be >= 1 (got %s)." % df
    return igamc(df/2, x/2)
 
def t_low(t, df):
    """Returns left-hand tail of Student's t distribution (-infinity to x).
    
    df, the degrees of freedom, ranges from 1 to infinity.
    Typically, df is (n-1) for a sample size of n.
    
    Result ranges from 0 to 1.
    
    See Cephes docs for details.
    """
    if df < 1:
        raise ValueError, "t_low: df must be >= 1 (got %s)." % df
    return stdtr(df, t)

def t_high(t, df):
    """Returns right-hand tail of Student's t distribution (x to infinity).
    
    df, the degrees of freedom, ranges from 1 to infinity.
    Typically, df is (n-1) for a sample size of n.
    
    Result ranges from 0 to 1.
    
    See Cephes docs for details.
    """
    if df < 1:
        raise ValueError, "t_high: df must be >= 1 (got %s)." % df
    return stdtr(df, -t) #distribution is symmetric

def tprob(t, df):
    """Returns both tails of t distribution (-infinity to -x, infinity to x)"""
    return 2 * t_high(abs(t), df)

def poisson_high(successes, mean):
    """Returns right tail of Poission distribution, Pr(X > x).
    
    successes ranges from 0 to infinity. mean must be positive.
    """
    return pdtrc(successes, mean)

def poisson_low(successes, mean):
    """Returns left tail of Poisson distribution, Pr(X <= x).

    successes ranges from 0 to infinity. mean must be positive.
    """
    return pdtr(successes, mean)

def poisson_exact(successes, mean):
    """Returns Poisson probablity for exactly Pr(X=successes).
    
    Formula is e^-(mean) * mean^(successes) / (successes)!
    """
    if successes == 0:
        return pdtr(0, mean)
    elif successes < mean:  #use left tail
            return pdtr(successes, mean) - pdtr(successes-1, mean)
    else: #successes > mean: use right tail
        return pdtrc(successes-1, mean) - pdtrc(successes, mean)

def binomial_high(successes, trials, prob):
    """Returns right-hand binomial tail (X > successes) given prob(success)."""
    if -1 <= successes < 0:
        return 1
    return bdtrc(successes, trials, prob)

def binomial_low(successes, trials, prob):
    """Returns left-hand binomial tail (X <= successes) given prob(success)."""
    return bdtr(successes, trials, prob)

def binomial_exact(successes, trials, prob):
    """Returns binomial probability of exactly X successes.
    
    Works for integer and floating point values.

    Note: this function is only a probability mass function for integer 
    values of 'trials' and 'successes', i.e. if you sum up non-integer 
    values you probably won't get a sum of 1.
    """
    if (prob < 0) or (prob > 1):
        raise ValueError, "Binomial prob must be between 0 and 1."
    if (successes < 0) or (trials < successes):
        raise ValueError, "Binomial successes must be between 0 and trials."
    return exp(ln_binomial(successes, trials, prob))

def f_low(df1, df2, x):
    """Returns left-hand tail of f distribution (0 to x).

    x ranges from 0 to infinity.
    
    Result ranges from 0 to 1.
    
    See Cephes docs for details.
    """
    return fdtr(df1, df2, x)
    
def f_high(df1, df2, x):
    """Returns right-hand tail of f distribution (x to infinity).
    
    Result ranges from 0 to 1.
    
    See Cephes docs for details.
    """
    return fdtrc(df1, df2, x)
  
def fprob(dfn, dfd, F, side='right'):
    """Returns both tails of F distribution (-inf to F and F to inf)

    Use in case of two-tailed test. Usually this method is called by
    f_two_sample, so you don't have to worry about choosing the right side.
    
    side: right means return twice the right-hand tail of the F-distribution.
        Use in case var(a) > var (b)
          left means return twice the left-hand tail of the F-distribution.
        Use in case var(a) < var(b)
    """
    if F < 0:
        raise ValueError, "fprob: F must be >= 0 (got %s)." % F
    if side=='right':
        return 2*f_high(dfn, dfd, F)
    elif side=='left':
        return 2*f_low(dfn, dfd, F)
    else:
        raise ValueError, "Not a valid value for side %s"%(side)

 
def stdtr(k, t):
    """Student's t distribution, -infinity to t.

    See Cephes docs for details.
    """
    if k <= 0:
        raise ValueError, 'stdtr: df must be > 0.'
    if t == 0:
        return 0.5
    if t < -2:
        rk = k
        z = rk / (rk + t * t)
        return 0.5 * betai(0.5 * rk, 0.5, z)
    #compute integral from -t to + t
    if t < 0:
        x = -t
    else:
        x = t

    rk = k  #degrees of freedom
    z = 1 + (x * x)/rk
    #test if k is odd or even
    if (k & 1) != 0:
        #odd k
        xsqk = x/sqrt(rk)
        p = atan(xsqk)
        if k > 1:
            f = 1
            tz = 1
            j = 3
            while (j <= (k-2)) and ((tz/f) > MACHEP):
                tz *= (j-1)/(z*j)
                f += tz
                j += 2
            p += f * xsqk/z
        p *= 2/PI
    else:
        #even k
        f = 1
        tz = 1
        j = 2
        while (j <= (k-2)) and ((tz/f) > MACHEP):
            tz *= (j-1)/(z*j)
            f += tz
            j += 2
        p = f * x/sqrt(z*rk)
    #common exit
    if t < 0:
        p = -p  #note destruction of relative accuracy
    p = 0.5 + 0.5 * p
    return p

def bdtr(k, n, p):
    """Binomial distribution, 0 through k.

    Uses formula bdtr(k, n, p) = betai(n-k, k+1, 1-p)

    See Cephes docs for details.
    """
    p = fix_rounding_error(p)
    if (p < 0) or (p > 1):
        raise ValueError, "Binomial p must be between 0 and 1."
    if (k < 0) or (n < k):
        raise ValueError, "Binomial k must be between 0 and n."
    if k == n:
        return 1
    dn = n - k
    if k == 0:
        return  pow(1-p, dn)
    else:
        return  betai(dn, k+1, 1-p)

def bdtrc(k, n, p):
    """Complement of binomial distribution, k+1 through n.

    Uses formula bdtrc(k, n, p) = betai(k+1, n-k, p)

    See Cephes docs for details.
    """
    p = fix_rounding_error(p)
    if (p < 0) or (p > 1):
        raise ValueError, "Binomial p must be between 0 and 1."
    if (k < 0) or (n < k):
        raise ValueError, "Binomial k must be between 0 and n."
    if k == n:
        return 0
    dn = n - k
    if k == 0:
        if p < .01:
            dk = -expm1(dn * log1p(-p))
        else:
            dk = 1 - pow(1.0-p, dn)
    else:
        dk = k + 1
        dk = betai(dk, dn, p)
    return dk
    
def pdtr(k, m):
    """Returns sum of left tail of Poisson distribution, 0 through k.

    See Cephes docs for details.
    """
    if k < 0:
        raise ValueError, "Poisson k must be >= 0."
    if m < 0:
        raise ValueError, "Poisson m must be >= 0."
    return igamc(k+1, m)

def pdtrc(k, m):
    """Returns sum of right tail of Poisson distribution, k+1 through infinity.

    See Cephes docs for details.
    """
    if k < 0:
        raise ValueError, "Poisson k must be >= 0."
    if m < 0:
        raise ValueError, "Poisson m must be >= 0."
    return igam(k+1, m)

def fdtr(a, b, x):
    """Returns left tail of F distribution, 0 to x.

    See Cephes docs for details.
    """
    if min(a, b) < 1:
        raise ValueError, "F a and b (degrees of freedom) must both be >= 1."
    if x < 0:
        raise ValueError, "F distribution value of f must be >= 0."
    w = a * x
    w /= float(b + w)
    return betai(0.5 * a, 0.5 * b, w)
    

def fdtrc(a, b, x):
    """Returns right tail of F distribution, x to infinity.

    See Cephes docs for details.
    """
    if min(a, b) < 1:
        raise ValueError, "F a and b (degrees of freedom) must both be >= 1."
    if x < 0:
        raise ValueError, "F distribution value of f must be >= 0."
    w = float(b) / (b + a * x)
    return betai(0.5 * b, 0.5 * a, w)

def gdtr(a, b, x):
    """Returns integral from 0 to x of Gamma distribution with params a and b.
    """
    if x < 0.0:
        raise ZeroDivisionError, "x must be at least 0."
    return igam( b, a * x)

def gdtrc(a, b, x):
    """Returns integral from x to inf of Gamma distribution with params a and b.
    """
    if x < 0.0:
        raise ZeroDivisionError, "x must be at least 0."
    return igamc(b, a * x)

#note: ndtri for the normal distribution is already imported

def chdtri(df, y):
    """Returns inverse of chi-squared distribution."""
    y = fix_rounding_error(y)
    if(y < 0.0 or y > 1.0 or df < 1.0):
        raise ZeroDivisionError, "y must be between 0 and 1; df >= 1"
    return 2 * igami(0.5*df, y)

def stdtri(k, p):
    """Returns inverse of Student's t distribution. k = df."""
    p = fix_rounding_error(p)
    # handle easy cases
    if k <= 0 or p < 0.0 or p > 1.0:
        raise ZeroDivisionError, "k must be >= 1, p between 1 and 0."
    rk = k
    #handle intermediate values
    if p > 0.25 and p < 0.75:
        if p == 0.5:
            return 0.0
        z = 1.0 - 2.0 * p;
        z = incbi(0.5, 0.5*rk, abs(z))
        t = sqrt(rk*z/(1.0-z))
        if p < 0.5:
            t = -t
        return t
    #handle extreme values
    rflg = -1
    if p >= 0.5:
            p = 1.0 - p;
            rflg = 1
    z = incbi(0.5*rk, 0.5, 2.0*p)

    if MAXNUM * z < rk:
        return rflg * MAXNUM
    t = sqrt(rk/z - rk)
    return rflg * t

def pdtri(k, p):
    """Inverse of Poisson distribution.

    Finds Poission mean such that integral from 0 to k is p.
    """
    p = fix_rounding_error(p)
    if k < 0 or p < 0.0 or p >= 1.0:
        raise ZeroDivisionError, "k must be >=0, p between 1 and 0."
    v = k+1;
    return igami(v, p)

def bdtri(k, n, y):
    """Inverse of binomial distribution.

    Finds binomial p such that sum of terms 0-k reaches cum probability y.
    """
    y = fix_rounding_error(y)
    if y < 0.0 or y > 1.0:
        raise ZeroDivisionError, "y must be between 1 and 0."
    if k < 0 or n <= k:
        raise ZeroDivisionError, "k must be between 0 and n"
    dn = n - k
    if k == 0:
        if y > 0.8:
            p = -expm1(log1p(y-1.0) / dn)
        else:
            p = 1.0 - y**(1.0/dn)
    else:
        dk = k + 1;
        p = incbet(dn, dk, 0.5)
        if p > 0.5:
            p = incbi(dk, dn, 1.0-y)
        else:
            p = 1.0 - incbi(dn, dk, y)
    return p

def gdtri(a, b, y):
    """Returns Gamma such that y is the probability in the integral.
    
    WARNING: if 1-y == 1, gives incorrect result. The scipy implementation
    gets around this by using cdflib, which is in Fortran. Until someone
    gets around to translating that, only use this function for values of
    p greater than 1e-15 or so!
    """
    y = fix_rounding_error(y)
    if y < 0.0 or y > 1.0 or a <= 0.0 or b < 0.0:
        raise ZeroDivisionError, "a and b must be non-negative, y from 0 to 1."
    return igami(b, 1.0-y) / a

def fdtri(a, b, y):
    """Returns inverse of F distribution."""
    y = fix_rounding_error(y)
    if( a < 1.0 or b < 1.0 or y <= 0.0 or y > 1.0):
        raise ZeroDivisionError, "y must be between 0 and 1; a and b >= 1"
    y = 1.0-y
    # Compute probability for x = 0.5
    w = incbet(0.5*b, 0.5*a, 0.5)
    # If that is greater than y, then the solution w < .5.
    # Otherwise, solve at 1-y to remove cancellation in (b - b*w).
    if w > y or y < 0.001:
            w = incbi(0.5*b, 0.5*a, y)
            x = (b - b*w)/(a*w)
    else:
            w = incbi(0.5*a, 0.5*b, 1.0-y)
            x = b*w/(a*(1.0-w))
    return x