# bnewbold_stats.py
# helper functions!
# @author: bryan newbold
# @email: bnewbold@mit.edu
# @date: oct 2006

import scipy, pylab

def floored_val(x,min=1.):
    if pylab.size(x) > 1:
        ret = scipy.array(pylab.copy.copy(x));
        for i in range(0,pylab.size(x)):
            ret[i] = floored_val(x[i],min=min);
        return ret;
    if(abs(x) < min):
        return min;
    else: 
        return abs(x);

def floored_root(x,min=1.):
    return scipy.sqrt(floored_val(x,min=min));

def rchisqr(tp,ffunc,tx,ty,terr):
    """finds reduced chi squared for a function compared to data"""
    
    if not (pylab.size(tx) == pylab.size(ty) == pylab.size(terr)):
        print "ERROR: input matrix mismatch... (rchisqr)";
        return;
    try:
        ffunc(tp,tx[-1]);
    except:
        print "problem with passed function and/or parameters!";
    try:
        eta_err = (ty - ffunc(tp,tx))**2 / (terr**2)
    except:
        print "i think you have a zero error in there...";
        eta_err = (ty - ffunc(tp,tx))**2 / (terr**2)
    return sum(eta_err)/(pylab.size(tx)-pylab.size(tp));


def rchisqr_poisson(tp,ffunc,tx,ty,floor=1.):
    """finds reduced chi squared for a function compared to data,
    assuming each data point has a poisson distribution and thus
    an error of sqrt(val). The error value is floored at 1.0 to 
    prevent stupid val=0 cases..."""
    
    if not (pylab.size(tx) == pylab.size(ty)):
        print "ERROR: input matrix mismatch... (rchisqr_poisson)";
        return;
    terr = floored_root(tx,min=floor);
    return rchisqr(tp,ffunc,tx,ty,terr);

def fittodata(tx,ty,p0,ffunc,err=None,floor=1.,ret_all=False,param_names=None):
    """does a fit and returns a dict of shit. 
    ret_all=True gives more shit."""
    if err == None:
        err = floored_root(ty,min=floor);
    if not (pylab.size(tx) == pylab.size(ty) == pylab.size(err)):
        print "ERROR: input matrix mismatch... (fittodata)";
        return;
    
    errfunc = lambda n,m,w: ffunc(n,m) - w;
   
    tx = scipy.array(tx);
    #tx = tx.tolist(); tx = scipy.array(tx);
    ty = scipy.array(ty);
    #ty = ty.tolist(); ty = scipy.array(ty);
    err = err.tolist(); err = scipy.array(err);
    p0 = scipy.array(p0).tolist(); p0 = scipy.array(p0);
    
    
    p1,cov_x,extra1,extra2,success = scipy.optimize.leastsq(errfunc, 
                    p0.copy(), args=(tx,ty), full_output=1)
    if size(p1) == 1:
        p1 = [p1,]
    print "DEBUG: done with fit"
    print "p1: " + str(p1)
    rcs = rchisqr(p1,ffunc,tx,ty,err);    
    try:
        p1_err = p1.copy()
    except:
        p1_err = [0,]
    for i in range(scipy.size(p1)):
        try:
            scipy.size(cov_x);
        except:
            p1_err = array([0.0,]) 
            p1 = array([p1,]) 
            break;
        if scipy.size(cov_x) == 1:
            p1_err = array([scipy.sqrt(cov_x),]);
            p1 = array([p1,]);
        else:
            p1_err[i] = scipy.sqrt(cov_x[i,i]);

    if ret_all:
        return {"params":p1, 
                "param_error":p1_err,
                "rchisqr":rcs,
                "ffunc":ffunc,
                "err":err,
                "x":tx,
                "y":ty,
                "success":success,
                "param_names":param_names};
    else:
        return {"params":p1,
                "param_error":p1_err,
                "rchisqr":rcs,
                "param_names":param_names};

def fittodata_gaussian(tx,ty,p0,err=None,floor=1.,ret_all=False):
    """runs fittodata for a gaussian (first parameter is mean,
    second is stddev)"""
    #gauss_func = lambda a,b: a[2]/(a[1] * sqrt(2*pi)) * exp(-.5 * ((b - a[0])/a[1])**2)
    #thing = std(ty);
    bigsum = scipy.integrate.simps(ty,x=tx)
    gauss_func = lambda a,b: bigsum/(a[1] * scipy.sqrt(2*pi)) * scipy.exp(-.5 * ((b - a[0])/a[1])**2)
    return fittodata(tx,ty,p0,gauss_func,err=err,floor=floor,ret_all=ret_all,
            param_names=["Mean","StdDev"]);

def fittodata_gaussian_amp(tx,ty,p0,err=None,floor=1.,ret_all=False):
    """runs fittodata for a gaussian with amplitude parameter (first parameter 
    is mean, second is stddev, third is amplitude)"""
    bigsum = scipy.integrate.simps(ty,x=tx)
    gauss_func = lambda a,b: a[2] * bigsum/(a[1] * scipy.sqrt(2*pi)) * scipy.exp(-.5 * ((b - a[0])/a[1])**2)
    return fittodata(tx,ty,p0,gauss_func,err=err,floor=floor,ret_all=ret_all,
            param_names=["Mean","StdDev","Amplitude"]);

def fittodata_gaussian_plus(tx,ty,p0,err=None,floor=1.,ret_all=False):
    """runs fittodata for a gaussian with extra parameters (first parameter 
    is mean, second is stddev, third is amplitude, fourth is constant offset, 
    fifth is linear trend)"""
    bigsum = scipy.integrate.simps(ty,x=tx)
    gauss_func = lambda a,b: a[2] * bigsum/(a[1] * scipy.sqrt(2*pi)) * scipy.exp(-.5 * ((b - a[0])/a[1])**2) + a[3] + a[4]*b
    return fittodata(tx,ty,p0,gauss_func,err=err,floor=floor,ret_all=ret_all,
            param_names=["Mean","StdDev","Amplitude","VertOffset","Slope"]);

def fittodata_linear(tx,ty,p0,err=None,floor=1.,ret_all=False):
    """runs fittodata for a linear curve (first param is slope,
    second is vertical offset)"""
    lin_func = lambda a,b: a[0] * b + a[1]
    return fittodata(tx,ty,p0,lin_func,err=err,floor=floor,ret_all=ret_all,
            param_names=["VertOffset (m)","Slope (b)"]);

def fittodata_quadratic(tx,ty,p0,err=None,floor=1.,ret_all=False):
    """runs fittodata for a quadratic curve (first param is a,
    second is b, third is c: ax^2+bx+c)"""
    quad_func= lambda a,b: a[0] * b*b  + a[1]*b + a[2]
    return fittodata(tx,ty,p0,quad_func,err=err,floor=floor,ret_all=ret_all,
            param_names=["QuadaticTerm (a)","VertOffset (b)","Slope (c)"]);

def fittodata_poisson(tx,ty,p0,err=None,floor=.5,ret_all=False):
    """runs fittodata for a poisson (first param is the mean)"""
    bigsum = scipy.integrate.simps(ty,x=tx)
    if type(p0) == type(0.1):
        p0 = [p0,]
    poissfunc = lambda a,b: bigsum * scipy.power(a[0],b) / scipy.factorial(b) * scipy.exp(-1 * a[0])
    return fittodata(tx,ty,p0,poissfunc,err=err,floor=floor,ret_all=ret_all,
            param_names=["Mean"]);

def fittodata_exp(tx,ty,p0,err=None,floor=1.,ret_all=False):
    """runs fittodata for an exponential (first param is x0,
    second param is tau, third param is vertical offset)"""
    exp_func = lambda a,b: a[0] * e ** (- b/a[1]) + a[2]
    return fittodata(tx,ty,p0,exp_func,err=err,floor=floor,ret_all=ret_all,
            param_names=["x0","tau","VertOffset"]);

def fittodata_lorentz(tx,ty,p0,err=None,floor=1.,ret_all=False):
    """runs fittodata for a lorentzian (first param is gamma,
    second param is mean)"""
    bigsum = scipy.integrate.simps(ty,x=tx)
    if type(p0) == type(0.1):
        p0 = [p0,]
    lorentz_func = lambda a,b: bigsum * a[0]/(2.*pi*((a[1]-b)**2+(a[0]/2.)**2. )) 
    return fittodata(tx,ty,p0,lorentz_func,err=err,floor=floor,ret_all=ret_all,
            param_names=["Gamma","Mean"]);

def fittodata_lorentz_plus(tx,ty,p0,err=None,floor=1.,ret_all=False):
    """runs fittodata for a lorentzian with extra parameters (first param 
    is gamma, second param is mean, third is constant offset, fourth is
    linear trend, fifth is amplitute)"""
    bigsum = scipy.integrate.simps(ty,x=tx)
    if type(p0) == type(0.1):
        p0 = [p0,]
    lorentz_func = lambda a,b,: a[4]*bigsum * a[0]/(2.*pi*((a[1]-b)**2+(a[0]/2.)**2. )) + a[2] + a[3]*b
    return fittodata(tx,ty,p0,lorentz_func,err=err,floor=floor,ret_all=ret_all,
            param_names=["Gamma","Mean","VertOffset","Slope","Amplitude"]);

def plotfit(indict,newfig=True):
    if not indict.has_key('x'):
        print "forgot to use ret_all on the fittodata?";
        return;
    p1 = indict['params'];
    x = indict['x']
    y = indict['y']
    err = indict['err']
    ffunc = indict['ffunc']
    param_names = indict['param_names']
    if newfig: pylab.figure();
    pylab.errorbar(x,y,fmt='b.',yerr=err)
    d=pylab.plot(x,y,'k.')
    fx = x
    if pylab.size(fx) < 100:
        fx = r_[min(x):max(x):(max(x)-min(x))/150.]
    f=pylab.plot(fx,ffunc(p1,fx),'r-')
    #fy = indict['ffunc'](indict['params'],indict['x'])
    pylab.title("Rough Fit of Data")
    pylab.xlabel("X")
    pylab.ylabel("Y")
    p1 = indict['params']
    p1_err = indict['param_error']
    
    ax = pylab.gca()
    extra_txt = 'reduced chi^2: %g'% (indict['rchisqr'])
    for i in range(scipy.size(p1)):
        if len(param_names)>i:
            extra_txt += '\n%s: %.3g +/- %.2g' % (param_names[i],p1[i],p1_err[i]);
        else:
            extra_txt += '\np[%i]: %.3g +/- %.2g' % (i,p1[i],p1_err[i]);
    t = pylab.text(.77, .75, extra_txt,
        horizontalalignment='center',
        verticalalignment='center',
        transform = ax.transAxes)
    pylab.legend((d,f),('Data', 'Fit'))