#!/usr/bin/env python
  
#  Copyright (c) 2005 Gavin E. Crooks <gec@threeplusone.com>
#
#  This software is distributed under the MIT Open Source License.
#  <http://www.opensource.org/licenses/mit-license.html>
#
#  Permission is hereby granted, free of charge, to any person obtaining a 
#  copy of this software and associated documentation files (the "Software"),
#  to deal in the Software without restriction, including without limitation
#  the rights to use, copy, modify, merge, publish, distribute, sublicense,
#  and/or sell copies of the Software, and to permit persons to whom the
#  Software is furnished to do so, subject to the following conditions:
#
#  The above copyright notice and this permission notice shall be included
#  in all copies or substantial portions of the Software.
#
#  THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 
#  IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 
#  FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
#  AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 
#  LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
#  OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN 
#  THE SOFTWARE.
#


""" Various bits of useful math not in the standard python library.

Constants :

- euler_gamma  = 0.577215...
- catalan      = 0.915965...
- golden_ratio = 1.618033...
- bits_per_nat = log2(e) = 1/log(2) 
- sqrt_2pi     = 2.50662...
    
Special Functions :



- gamma()                       -- Gamma function.
- lngamma()                     -- Logarithm of the gamma function
- factorial()                   -- Factorial function.
- digamma()                     -- Digamma function (logarithmic derivative of gamma).
- trigamma()                    -- Trigamma function (derivative of digamma).

- cgamma()                       -- Complex math version of counterparts above.
- clngamma()                                   
- cdigamma()                     
- ctrigamma()                    

- entropy()                     -- The entropy of a probability vector
- incomplete_gamma()            -- The 'upper' incomplete gamma function.
- normalized_incomplete_gamma() -- 
- log2()                          -- Base 2 logarithms.
- argmin()
- argmax()


"""
from __future__ import division


__all__ = ('euler_gamma', 'catalan', 'golden_ratio', 'bits_per_nat', 'sqrt_2pi',
            'gamma', 'lngamma', 'factorial', 'digamma', 'trigamma',
            'cgamma', 'clngamma', 'cdigamma', 'ctrigamma',            
            'entropy', 'log2',            
            'incomplete_gamma', 'normalized_incomplete_gamma', 
            'argmax', 'argmin'
            )            

from math import *
import cmath as cm

from itertools import count

from ._py3k import zip


# Some mathematical constants
euler_gamma  = 0.57721566490153286060651
catalan      = 0.91596559417721901505460
golden_ratio = 1.6180339887498948482046
bits_per_nat = 1.44269504088896340735992468100 # = log_2(e) = 1/log(2) 
sqrt_2pi     = 2.5066282746310005024157652848110





# The Lanczos approximation for the gamma function is
#                          
#                -(z + g + 1/2)               (z + 1/2)                   
# Gamma(z+1) =  e               * (z + g + 1/2)        * Sqrt(2Pi) * C 
#                                                                        
#
#                      c[1]    c[2]    c[3]
#          C = [c[0] + ----- + ----- + ----- + ...   ]
#                      z + 1   z + 2   z + 3 
#
#
#  To calculate digamma and trigamma functions we take an analytic derivative
#  of the Lanczos approximation.
#
#  Gamma(z)  = Gamma(z+1)/z
#  Digamma(z) = D ln Gamma(z)
#  Trigamma(z) = D Digamma(z)         

# These Lanczos constants are from
# "A note on the computation of the convergent
# Lanczos complex Gamma approximation." Paul Godfrey (2001)
# http://my.fit.edu/~gabdo/gamma.txt


__lanczos_gamma = 607./128.
__lanczos_coefficients = ( 
       0.99999999999999709182,
      57.156235665862923517,
     -59.597960355475491248,
      14.136097974741747174,
      -0.49191381609762019978,
        .33994649984811888699e-4,
        .46523628927048575665e-4,
       -.98374475304879564677e-4,
        .15808870322491248884e-3,
       -.21026444172410488319e-3,
        .21743961811521264320e-3,
       -.16431810653676389022e-3,
        .84418223983852743293e-4,
       -.26190838401581408670e-4,
        .36899182659531622704e-5)

__factorial =(
    1.,
    1.,
    2.,
    6.,
    24.,
    120.,
    720.,
    5040.,
    40320.,
    362880.,
    3628800.,
    39916800.,
    479001600.,
    6227020800.,
    87178291200.,
    1307674368000.,
    20922789888000.,
    355687428096000.,
    6402373705728000.,
    121645100408832000.,
    2432902008176640000.,
    51090942171709440000.,
    1124000727777607680000.,
    25852016738884976640000.,
    620448401733239439360000.,
    15511210043330985984000000.,
    403291461126605635584000000.,
    10888869450418352160768000000.,
    304888344611713860501504000000.,
    8841761993739701954543616000000.,
    265252859812191058636308480000000.,
    8222838654177922817725562880000000.,
    263130836933693530167218012160000000. )
    
def gamma(x) :
    """The gamma function. Returns exact results for small integers. Will
    overflow for modest sized arguments. Use lngamma(z) instead. 
    
    See: Eric W. Weisstein. "Gamma Function." From MathWorld, A Wolfram Web Resource.
         http://mathworld.wolfram.com/GammaFunction.html

    """
    return cgamma(x).real


def cgamma(z) :
    """Gamma function with complex arguments and results."""
    
    z = complex(z)
    n = floor(z.real)   
    if  n == z  :
        if n <= 0 :
            return complex(1.0/0.0) # Infinity
        elif n <= len(__factorial) :
            return complex(__factorial[int(n)-1])
        
    zz = z
    if z.real < 0.5 :   # Reflection formula
        zz = 1-z
    
        
    g = __lanczos_gamma
    c = __lanczos_coefficients 
    
    zz = zz - 1.    
    zh = zz + 0.5
    zgh = zh + g
    zp = zgh** (zh*0.5) # trick for avoiding FP overflow above z=141

    ss = 0.0
    for k in range(len(c)-1,0,-1):
        ss += c[k]/(zz+k)
    
    f = (sqrt_2pi*(c[0]+ss)) * (( zp*cm.exp(-zgh)) *zp)

    if z.real<0.5 :
        f  = pi /( cm.sin(pi*z) *f)

    return complex(f)

def lngamma(x) :
    """The logarithm of the gamma function. 
    """
    return clngamma(x).real

   
def clngamma(z) :
    """The logarithm of the gamma function for a complex argument
    """
    z = complex(z)
    
    # common case optimization
    n = floor(z.real)   
    if  n == z  :
        if n <= 0 :
            return complex(1.0/0.0) # Infinity
        elif n <= len(__factorial) :
            return complex(__factorial[int(n)-1])
        
    zz = z
    if z.real < 0.5 :
        zz = 1-z
    
        
    g = __lanczos_gamma
    c = __lanczos_coefficients 
    
    zz = zz - 1.    
    zh = zz + 0.5
    zgh = zh + g
    zp = zgh** (zh*0.5) # trick for avoiding FP overflow above z=141

    ss = 0.0
    for k in range(len(c)-1,0,-1):
        ss += c[k]/(zz+k)
    
    f = (sqrt_2pi*(c[0]+ss)) * (( zp*cm.exp(-zgh)) *zp)

    if z.real<0.5 :
        f  = pi /( cm.sin(pi*z) *f)

    return cm.log(f)
          
                    
def factorial(z) :
    """ The factorial function. 
    factorial(z) == gamma(z+1)
    """
    return gamma(z+1)
              
def digamma(x) :   
    """The digamma function, the logarithmic derivative of the gamma function.
             digamma(z) = d/dz ln( gamma(z) )


     See: Eric W. Weisstein. "Digamma Function." From MathWorld--
     A Wolfram Web Resource. http://mathworld.wolfram.com/DigammaFunction.html
    """
    return cdigamma(x).real
          
          
def cdigamma(z) :
    """Digamma function with complex arguments"""
    z = complex(z)

    g = __lanczos_gamma
    c = __lanczos_coefficients 
        
    zz = z
    if z.real < 0.5 :
        zz = 1 -z
        
    n=0.
    d=0.
    for k in range(len(c)-1,0,-1):
        dz =1./(zz+(k+1)-2);
        dd =c[k] * dz
        d = d + dd 
        n = n - dd * dz

    d = d + c[0]
    gg = zz + g - 0.5
    f = cm.log(gg) + (n/d - g/gg)

    if z.real<0.5 :
        f -= pi / cm.tan( pi * z)
        
    return f


def trigamma(x) :
    """The trigamma function, the derivative of the digamma function.
            trigamma(z) = d/dz digamma(z) = d/dz d/dz ln( gamma(z) )
    
    See: Eric W. Weisstein. "Trigamma Function." From MathWorld--
    A Wolfram Web Resource. http://mathworld.wolfram.com/TrigammaFunction.html
    """
    return ctrigamma(x).real    

def ctrigamma(z):
    "The trigamma function with complex arguments and return value."
    z = complex(z)
    g = __lanczos_gamma
    c = __lanczos_coefficients 
        
    t1=0.
    t2=0.
    t3=0.
    for k in range(len(c)-1,0,-1):
        dz =1./(z+k);       
        dd1 = c[k]* dz
        t1 += dd1
        dd2 = dd1 * dz
        t2 += dd2
        t3 += dd2 * dz

    t1 += c[0]
    c =  - (t2*t2)/(t1*t1)  +2*t3/t1

    result = 1./(z*z)
    gg = z + g + 0.5
    result += - (z+0.5)/ (gg*gg)
    result += 2./gg

    result += c

    return result

def incomplete_gamma(a,x) :
    """The 'upper' incomplete gamma function:

                            oo
                             -
                            |    -t  a-1
    incomplete_gamma(a,x) = |   e   t   dt.
                            |
                           -                         
                            x

    In Mathematica, Gamma[a,x].

    Note that, very confusingly, the phrase 'incomplete gamma fucntion'
    can also refer to the same integral between 0 and x, (the 'lower'
    incomplete gamma function) or to the normalized versions,
    normalized_incomplete_gamma() )


    See: Eric W. Weisstein. "Gamma Function." From MathWorld, A Wolfram Web Resource.
         http://mathworld.wolfram.com/IncompleteGammaFunction.html

    Bugs :
        This implementation is not very accurate for some arguments. 
    """   
    return  normalized_incomplete_gamma(a,x) * gamma(a)

    
def normalized_incomplete_gamma(a,x) :
    """The upper, incomplete gamma function normalized so that the limiting
    values are zero and one.
    
     Q(a,x) = incomplete_gamma(a,x) / gamma(a) 

    See: 
        incomplete_gamma()
    Bugs :
        This implementation is not very accurate for some arguments. 
    """
    maxiter = 100
    epsilon = 1.48e-8
    small = 1e-30
    
        
    if a<=0 or x<0 : 
        raise ValueError("Invalid arguments")
    if x == 0.0 : return 1.0
    
    if x<= a+1 :
        # Use the series representation
        term = 1./a
        total = term
        for n in range(1,maxiter) :
            term *= x/(a+n)
            total += term
            if abs(term/total) < epsilon : 
                return 1. - total * exp(-x+a*log(x) - lngamma(a) )
        raise RuntimeError(
            "Failed to converge after %d iterations." % (maxiter) )
    else :
        # Use the continued fraction representation
        total = 1.0
        b = x + 1. -a
        c = 1./small
        d = 1./b
        h = d
        for i in range(1, maxiter) :
            an = -i * (i-a)
            b = b+2.
            d = an * d + b
            if abs(d) < small : d = small
            c = b + an /c
            if abs(c) < small : c= small
            d = 1./d
            term = d * c
            h = h * term
            if abs( term-1.) < epsilon :
                return h * exp(-x+a*log(x) - lngamma(a) )
        raise RuntimeError(
            "Failed to converge after %d iterations." % (maxiter) )


 
def log2( x) :
    """ Return the base 2 logarithm of x """
    return log(x,2)


def entropy( pvec, base= exp(1) ) :
    """ The entropy S = -Sum_i p_i ln p_i
        pvec is a frequency vector, not necessarily normalized. 
    """
    # TODO: Optimize
    if len(pvec) ==0 : 
        raise ValueError("Zero length vector")
        
        
    total = 0.0
    ent = 0.0
    for p in pvec:
        if p>0 : # 0 log(0) =0 
            total += p
            ent += - log(float(p)) *p
        elif p<0:
            raise ValueError("Negative probability")
    
    
    ent = (ent/total) + log(total)
    ent /= log(base)

    return ent


def argmax( alist) :
    """Return the index of the last occurrence of the maximum value in the list."""
    return max(zip(alist, count()))[1]


def argmin( alist) :
    """Return the index of the first occurrence of the minimum value in the list."""
    return min(zip(alist, count()))[1]