import argparse
import sys,os,glob
import math, time
pandapath=os.environ.get('PANDAPATH')
sys.path.append(pandapath+'/macro/tpc/FOPI/mberger')
from functions import *
from anaFile import anaFile

parser=argparse.ArgumentParser(description='Slice dEdx vs mom histogram, fit proton line with langaus and extract a TCut')
parser.add_argument("infile",help='the output file from analyseFOPIReco_pions.py',type=str)
parser.add_argument("fitinfo",help='the file with the start parameters for the fit',type=str)
parser.add_argument('--slicewidth',help='the width of the momentum bins (%(default)s)',type=float,default=-1)
parser.add_argument('--nslices',help='number of slices in momentum (%(default)s)',type=int,default=-1)
parser.add_argument('--momRange',help='the range of the momentum to slice (%(default)s)',type=float,nargs=2,default=[0,5])

args=parser.parse_args()

if args.slicewidth==-1 and args.nslices==-1:
    print "please give either slicewidth or nslices!"
    exit()
    
if args.slicewidth==-1:
    args.slicewidth=float(args.momRange[1]-args.momRange[0])/float(args.nslices)
if args.nslices==-1:
    args.nslices=int(float(args.momRange[1]-args.momRange[0])/float(args.slicewidth))
    
#read in fit start parameters
fitpar=[]
nlines=0
for line in open(args.fitinfo,"r"):
    if line[0]=='#':
        continue    
    nlines+=1
    words=line.split()
    pars=[float(word) for word in words]
    fitpar.append({})
    fitpar[-1]["start dedx"] =pars[0]
    fitpar[-1]["stop dedx"]  =pars[1]
    fitpar[-1]["mean gauss"] =pars[2]
    fitpar[-1]["sig gauss"]   =pars[3]
    fitpar[-1]["mean landau"]=pars[4]
    fitpar[-1]["sig landau"] =pars[5]
    
if nlines<args.nslices:
    print "too few parameters given. neede {0}, got {1}".format(args.nslices,nlines)
    exit()
    
if os.path.isdir(args.infile):
    pass
else:
    anafile=anaFile(args.infile)

cdEdx=ROOT.TCanvas("cdEdx","dEdx",700,1000)
cdEdx.Divide(1,2)
cSlices=ROOT.TCanvas("cslices","slices",1500,1500)
cSlices.Divide(int(math.ceil(math.sqrt(args.nslices))),int(math.ceil(math.sqrt(args.nslices))))
cFit=ROOT.TCanvas("cFit","fit",700,500)

hdedx=anafile.getHist("hdedxVsmom")

cdEdx.cd(1)
hdedx.SetStats(0)
hdedx.Draw("colz")


#slice the dedx vs mom spectrum in momentum slices
print 'looping over',args.nslices,'with a width of',args.slicewidth
start=args.momRange[0]
stop=start+args.slicewidth
momslices=[]
cutlines=[]
islice=-1
while start<args.momRange[1]:
    islice+=1
    start_bin=hdedx.GetXaxis().FindBin(start)
    stop_bin =hdedx.GetXaxis().FindBin(stop)
    dedx_max=hdedx.GetYaxis().GetXmax()
    print "making line from ({0},{1})  to ({2},{3})".format(start,dedx_max,stop,dedx_max)
    cutlines.append(ROOT.TLine(start,0,start,dedx_max))
    cutlines[-1].SetLineColor(ROOT.kPink+10)
    cdEdx.cd(1)
    cutlines[-1].Draw("same")
    print "projecting from bin {0} to bin {1}".format(start_bin,stop_bin)
    momslices.append(hdedx.ProjectionY("{0}_py_{1}({2}-{3})".format(hdedx.GetName(),islice,start,stop),start_bin,stop_bin))
    momslices[-1].SetTitle(momslices[-1].GetName())
    cSlices.cd(islice+1)
    momslices[-1].Draw()
    start=stop
    stop+=args.slicewidth
    

cdEdx.Update()

#perform the fitting
fitcollection=[]
islice=-1
for slice in momslices:
    unhappy=True
    mg_start=200
    sg_start=20
    ml_start=200
    sl_start=20
    while unhappy:
        islice+=1
        fitcollection.append({})
        dedx=ROOT.RooRealVar("dE/dx [a.u.]","dEdX [a.u.]",fitpar[islice]["start dedx"],fitpar[islice]["stop dedx"]);     
        dedx_set=ROOT.RooArgList("dedx_set")
        dedx_set.add(dedx)
        fitcollection[-1]["data"]=ROOT.RooDataHist("momslice_{0}".format(islice),"momslice_{0}".format(islice),dedx_set,slice)
        
        frame=dedx.frame(ROOT.RooFit.Title("{0}".format(slice.GetTitle())))
         #construct fit function
        print 'fit parameters for this slice:',fitpar[islice]
         #gauss
        
        fitcollection[-1]["mean_gaus"]=ROOT.RooRealVar("mean_gaus{0}".format(islice),
                                                       "mean_gaus{0}".format(islice),
                                                       fitpar[islice]['mean gauss'],
                                                       fitpar[islice]['mean gauss']-10*fitpar[islice]['sig gauss'],
                                                       fitpar[islice]['mean gauss']+10*fitpar[islice]['sig gauss'])
        fitcollection[-1]["sigma_gaus"]=ROOT.RooRealVar("sigma_gaus{0}".format(islice),
                                                        "sigma_gaus{0}".format(islice),
                                                        fitpar[islice]['sig gauss'],
                                                        0,2*fitpar[islice]['sig gauss'])
        fitcollection[-1]["gauss"]=ROOT.RooGaussian("gauss{0}".format(islice),"gauss{0}".format(islice),dedx,fitcollection[-1]["mean_gaus"],fitcollection[-1]["sigma_gaus"])
    
        #landau
        fitcollection[-1]["mean_landau"]=ROOT.RooRealVar("mean_landau{0}".format(islice),
                                                         "mean_landau{0}".format(islice),
                                                         fitpar[islice]['mean landau'],
                                                         fitpar[islice]['mean landau']-10*fitpar[islice]['sig landau'],
                                                         fitpar[islice]['mean landau']+10*fitpar[islice]['sig landau'])
        fitcollection[-1]["sigma_landau"]=ROOT.RooRealVar("sigma_landau{0}".format(islice),
                                                          "sigma_landau{0}".format(islice),
                                                          fitpar[islice]['sig landau'],
                                                          0,2*fitpar[islice]['sig landau'])
        fitcollection[-1]["landau"]=ROOT.RooLandau("landau{0}".format(islice),"landau{0}".format(islice),dedx,fitcollection[-1]["mean_landau"],fitcollection[-1]["sigma_landau"])
        
        #Convolution of landau and gauss
        fitcollection[-1]["langau"]=ROOT.RooFFTConvPdf("langau{0}".format(islice),"landau (X) gauss {0}".format(islice),dedx,fitcollection[-1]["landau"],fitcollection[-1]["gauss"]); 
    
        fitcollection[-1]["langau"].fitTo(fitcollection[-1]["data"])
        #fitcollection[-1]['gauss'].fitTo(fitcollection[-1]["data"])
        #fitcollection[-1]['landau'].fitTo(fitcollection[-1]["data"])
        
        fitcollection[-1]["data"].plotOn(frame)
        fitcollection[-1]["data"].statOn(frame)
        fitcollection[-1]["langau"].plotOn(frame,ROOT.RooFit.Range("Full"))
        fitcollection[-1]["langau"].paramOn(frame)
        #fitcollection[-1]['gauss'].plotOn(frame,ROOT.RooFit.LineColor(ROOT.kRed))
        #fitcollection[-1]['gauss'].paramOn(frame)
        #fitcollection[-1]['landau'].plotOn(frame,ROOT.RooFit.LineColor(ROOT.kGreen))
        #fitcollection[-1]['landau'].paramOn(frame)
        
        
        
        cFit.cd()
        frame.Draw()
        cFit.Update()
        happy=raw_input("is it good? (y/n)")
        if happy=="y":
            unhappy=False
        else:
            mg_start=float(raw_input("mean for gauss"))
            sg_start=float(raw_input("sigma for gauss"))
            ml_start=float(raw_input("mean for landau"))
            sl_start=float(raw_input("sigma for landau"))

thisIsTheEnd()