"""Python functions from Numerical Recipes

Functions:
avevar    Mean and variance of a list of numbers.
factrl    Factorial.
bico      Binomial coefficient.
chsone    Chi-square test between observed and expected events.
chstwo    Chi-square test for two sets of binned data.
frprmn    Minimize a function with conjugate gradient descent.

"""
import nr

def _pylist2floatArray(l, startindex=0):
    # Turn a python list into a C array.
    # borrowed heavily from the SWIG manual
    a = nr.new_floatArray(len(l)+startindex)
    for i in range(len(l)):
        nr.floatArray_setitem(a, i+startindex, l[i])
    return a

def _floatArray2pylist(a, nitems, startindex=0):
    # Turn a C array into a python list.
    l = []
    for i in range(nitems):
        l.append(nr.floatArray_getitem(a, i+startindex))
    return l

def avevar(data):
    """avevar(data) -> ave, var"""
    c_data = _pylist2floatArray(data, 1)
    ave, var = nr.avevar(c_data, len(data))
    nr.delete_floatArray(c_data)
    return ave, var

def factrl(n):
    """factrl(n) -> n!"""
    f = nr.factrl(n)
    if nr.cvar.ERROR:
        error = nr.cvar.ERROR
        nr.clear_error()
        raise ValueError, error
    return f
    
def bico(n, k):
    """bico(n, k) -> n choose k"""
    f = nr.bico(n, k)
    if nr.cvar.ERROR:
        error = nr.cvar.ERROR
        nr.clear_error()
        raise ValueError, error
    return f

def chsone(bins, ebins, knstrn):
    """chsone(bins, ebins, knstrn) -> df, chsq, prob"""
    if len(bins) != len(ebins):
        raise ValueError, "bins and ebins should be same length"
    c_bins = _pylist2floatArray(bins, 1)
    c_ebins = _pylist2floatArray(ebins, 1)
    df, chsq, prob = nr.chsone(c_bins, c_ebins, len(bins), knstrn)
    nr.delete_floatArray(c_bins)
    nr.delete_floatArray(c_ebins)
    if nr.cvar.ERROR:
        error = nr.cvar.ERROR
        nr.clear_error()
        raise ValueError, error
    return df, chsq, prob

def chstwo(bins1, bins2, knstrn):
    """chstwo(bins1, bins2, knstrn) -> df, chsq, prob"""
    if len(bins1) != len(bins2):
        raise ValueError, "bins1 and bins2 should be same length"
    c_bins1 = _pylist2floatArray(bins1, 1)
    c_bins2 = _pylist2floatArray(bins2, 1)
    df, chsq, prob = nr.chstwo(c_bins1, c_bins2, len(bins1), knstrn)
    nr.delete_floatArray(c_bins1)
    nr.delete_floatArray(c_bins2)
    if nr.cvar.ERROR:
        error = nr.cvar.ERROR
        nr.clear_error()
        raise ValueError, error
    return df, chsq, prob

def frprmn(p, ftol, func, dfunc):
    """frprmn(p, ftol, func, dfunc) -> p, num_iterations, lowest value

    Use conjugate gradient descent to minimize a function.  p is the
    initial starting point.  ftol is the tolerance used to determine
    when to stop.  1E-4 is a reasonable value.  func and dfunc are
    callback functions.  func takes a point and evaluates the function
    at that point.  dfunc takes a point and returns a list of the
    partial derivatives evaluated at that point.

    """
    c_p = _pylist2floatArray(p, 1)
    x = nr.frprmnpy(c_p, len(p), ftol, func, dfunc)
    iter, fret = x
    p = _floatArray2pylist(c_p, len(p), 1)
    nr.delete_floatArray(c_p)
    if nr.cvar.ERROR:
        error = nr.cvar.ERROR
        nr.clear_error()
        raise AssertionError, error
    return p, iter, fret