from pylab import *
import pylab
from scipy import *

def plotfit_gaus(indict,newfig=True,txtp=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']
    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("Gaussian Fit of Events")
    pylab.xlabel("Number of Charged Particles")
    pylab.ylabel("Events")
    p1 = indict['params']
    p1_err = indict['param_error']

    ax = pylab.gca()
#    ax.set_xlim([-1,14])
    extra_txt = 'reduced chi squared: %g'\
        % (indict['rchisqr'])
    for i in range(scipy.size(p1)):
        extra_txt += '\np[%i]: %g +/- %g' \
            % (i,p1[i],p1_err[i]);

#    if(pylab.size(p1) == 2):
#        extra_txt = 'reduced chi squared: %g\np[0]: %g\np[1]: %g' \
#            % (indict['rchisqr'],p1[0],p1[1])
#    else:
#        extra_txt= 'reduced chi squared:%g\np[0]:%g, +/- %g\n p[1]:%g\np[2]:%g' \
#            % (indict['rchisqr'],p1[0],p1[1],p1[2])

    if txtp:
        t = pylab.text(.77, .75, extra_txt,
            horizontalalignment='center',
            verticalalignment='center',
            transform = ax.transAxes)
        pylab.legend((d,f),('Data', 'Fit'))

def plotfit_poiss(indict,newfig=True,txtp=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']
#    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.]
    p1= p1[0]
    if txtp:
        f2=pylab.plot(fx,ffunc(p1,fx),'r-')
    else:
        f2=pylab.plot(fx,ffunc(p1,fx),'g-')
    #fy = indict['ffunc'](indict['params'],indict['x'])
    pylab.title("Poisson Fit of Events")
    pylab.xlabel("Number of Charged Particles")
    pylab.ylabel("Events")
    p1 = indict['params']
    p1_err = indict['param_error']

    p1= p1[0]
    ax = pylab.gca()
    #ax.set_xlim([-1,20])
    extra_txt = 'reduced chi squared: %g'\
        % (indict['rchisqr'])
    for i in range(scipy.size(p1)):
        if type(p1_err) is not type(9.2):
            extra_txt += '\np[%i]: %g +/- %g' \
                % (i,p1[i],p1_err[i]);
        else:
            extra_txt += '\np[%i]: %g' \
                % (i,p1[i]);

#    if(pylab.size(p1) == 2):
#        extra_txt = 'reduced chi squared: %g\np[0]: %g\np[1]: %g' \
#            % (indict['rchisqr'],p1[0],p1[1])
#    else:
#        extra_txt= 'reduced chi squared:%g\np[0]:%g, +/- %g\n p[1]:%g\np[2]:%g' \
#            % (indict['rchisqr'],p1[0],p1[1],p1[2])
    if txtp:
        t = pylab.text(.77, .75, extra_txt,
            horizontalalignment='center',
            verticalalignment='center',
            transform = ax.transAxes)
    if txtp:
        pylab.legend((d,f2),('Data', 'Fit'))

lsd = pylab.load("lineshapedata.txt")
ls1 = pylab.load("lineshape1.txt")

rawlsd = []
for i in range(size(lsd[0,:])):
    for j in range(lsd[1,i]):
                rawlsd.append(lsd[0,i])

rawls1 = []
for i in range(size(ls1[0,:])):
    for j in range(ls1[1,i]):
                rawls1.append(ls1[0,i])

plotfit(fittodata_lorentz(lsd[0,:],lsd[1,:],[50.0,20.0],err=sqrt(lsd[1,:]),ret_all=True))
hist(rawlsd, bins=lsd[0,:],align='center',facecolor="lightblue")
pylab.title("Lorentzian fit to lineshapedata.txt")
plotfit(fittodata_gaussian(lsd[0,:],lsd[1,:],[50.0,20.0],err=sqrt(lsd[1,:]),ret_all=True))
hist(rawlsd, bins=lsd[0,:],align='center',facecolor="lightblue")
pylab.title("Gaussian fit to lineshapedata.txt")

plotfit(fittodata_lorentz(ls1[0,:],ls1[1,:],[50.0,20.0],err=sqrt(ls1[1,:]),ret_all=True))
hist(rawls1, bins=ls1[0,:],align='center',facecolor="lightblue")
pylab.title("Lorentzian fit to lineshape1.txt")
plotfit(fittodata_gaussian(ls1[0,:],ls1[1,:],[50.0,20.0],err=sqrt(ls1[1,:]),ret_all=True))
hist(rawls1, bins=ls1[0,:],align='center',facecolor="lightblue")
pylab.title("Gaussian fit to lineshape1.txt")

#pylab.figure()
#plot(ls1[0,:],ls1[1,:])