import ROOT,argparse

parser=argparse.ArgumentParser(description='plot fit comparison related stuff')
parser.add_argument('infile',help='the input file with the Result Tree',type=str)
parser.add_argument('--save',help='save the canvases to specified location',type=str,default='Nope')

args=parser.parse_args()

rfile=ROOT.TFile.Open(args.infile,'read')
tree = rfile.Get('Result')


meanU={}
meanU['mult']=ROOT.TH1D('hmeanUmult','mean U mult',120,-0.6,0.6)
meanU['2D']=ROOT.TH1D('hmeanU2D','mean U 2D',120,-0.6,0.6)
meanU['MC']=ROOT.TH1D('hmeanUMC','mean U MC',120,-0.6,0.6)

meanV={}
meanV['mult']=ROOT.TH1D('hmeanVmult','mean V mult',120,-0.6,0.6)
meanV['2D']=ROOT.TH1D('hmeanV2D','mean V 2D',120,-0.6,0.6)
meanV['MC']=ROOT.TH1D('hmeanVMC','mean V MC',120,-0.6,0.6)

fittime={}
fittime['mult']=ROOT.TH1D('hftimemult','RooFit function',75,0,15)
fittime['2D']=ROOT.TH1D('hftime2D','Free function',75,0,15)
fittime['diff']=ROOT.TH1D('hftimeDiff','Differecnce in fitting times (2D-Mult)',75,0,15)

failed={}
failed['2D']=0
failed['mult']=0

meanErr={}
meanErr['U']={}
meanErr['V']={}
meanErr['U']['2D']=ROOT.TH1D('hmeanErrU2D','Error Mean U 2D',100,0,1)
meanErr['V']['2D']=ROOT.TH1D('hmeanErrV2D','Error Mean V 2D',100,0,1)
meanErr['U']['mult']=ROOT.TH1D('hmeanErrUmult','Error Mean U mult',100,0,1)
meanErr['V']['mult']=ROOT.TH1D('hmeanErrVmult','Error Mean V mult',100,0,1)

for e in tree:
    efailed=False

    if e.qualMult!=-1 or e.statusMult!=0:
        failed['mult']+=1
        efailed=True
    if e.qual2D!=-1 or e.status2D!=0:
        failed['2D']+=1
        efailed=True

    if efailed:
        continue

    meanU['mult'].Fill(e.meanUMult)
    meanU['2D'].Fill(e.meanU2D)
    meanU['MC'].Fill(e.mcU)

    meanV['mult'].Fill(e.meanVMult)
    meanV['2D'].Fill(e.meanV2D)
    meanV['MC'].Fill(e.mcV)

    fittime['mult'].Fill(e.ftimeMult*100)
    fittime['2D'].Fill(e.ftime2D)
    fittime['diff'].Fill(e.ftime2D-e.ftimeMult)

    meanErr['U']['2D'].Fill(e.meanErrU2D)
    meanErr['V']['2D'].Fill(e.meanErrV2D)
    meanErr['U']['mult'].Fill(e.meanErrUMult)
    meanErr['V']['mult'].Fill(e.meanErrVMult)

labeloff=0.015
ytitoff=1.5
xtitoff=1.3

canv=[]
canv.append(ROOT.TCanvas('canv1','meanU'))
leg1=ROOT.TLegend(0.6,0.7,0.9,0.89)
leg1.SetBorderSize(0)
leg1.SetFillColor(0)
leg1.SetFillStyle(4000)
leg1.AddEntry(meanU['MC'],'Monte Carlo','l')
leg1.AddEntry(meanU['mult'],'RooFit function','l')
leg1.AddEntry(meanU['2D'],'Free function','l')


for h in meanU:
    meanU[h].SetStats(0)
    meanU[h].SetLineWidth(2)
    meanU[h].SetTitle('')
    meanU[h].GetXaxis().SetRangeUser(-0.4,0.4)


meanU['MC'].SetLineColor(ROOT.kAzure-2)
meanU['MC'].Draw('same')
meanU['MC'].SetMaximum(meanU['MC'].GetBinContent(meanU['MC'].GetMaximumBin())*1.1)
meanU['MC'].GetXaxis().SetTitle('U (cm)')
meanU['MC'].GetXaxis().SetLabelOffset(labeloff)
meanU['MC'].GetYaxis().SetLabelOffset(labeloff)
meanU['MC'].GetYaxis().SetTitle('Counts')
meanU['MC'].GetYaxis().SetTitleOffset(ytitoff)
meanU['MC'].GetXaxis().SetTitleOffset(xtitoff)

meanU['mult'].SetLineColor(ROOT.kSpring-1)
meanU['mult'].Draw("same")
meanU['2D'].SetLineColor(ROOT.kMagenta-3)
meanU['2D'].Draw('same')
leg1.Draw()


canv.append(ROOT.TCanvas('canv2','meanV'))
leg2=ROOT.TLegend(0.6,0.7,0.9,0.89)
leg2.SetBorderSize(0)
leg2.SetFillColor(0)
leg2.SetFillStyle(4000)
leg2.AddEntry(meanV['MC'],'Monte Carlo','l')
leg2.AddEntry(meanV['mult'],'RooFit function','l')
leg2.AddEntry(meanV['2D'],'Free function','l')


for h in meanV:
    meanV[h].SetStats(0)
    meanV[h].SetLineWidth(2)
    meanV[h].SetTitle('')
    meanV[h].GetXaxis().SetRangeUser(-0.4,0.4)
    meanV[h].GetXaxis().SetLabelOffset(labeloff)
    meanV[h].GetYaxis().SetLabelOffset(labeloff)

meanV['MC'].SetLineColor(ROOT.kAzure-2)
meanV['MC'].Draw('same')
meanV['MC'].SetMaximum(meanV['MC'].GetBinContent(meanV['MC'].GetMaximumBin())*1.1)
meanV['MC'].GetXaxis().SetTitle('V (cm)')
meanV['MC'].GetYaxis().SetTitle('Counts')
meanV['MC'].GetYaxis().SetTitleOffset(ytitoff)
meanV['MC'].GetXaxis().SetTitleOffset(xtitoff)

meanV['mult'].SetLineColor(ROOT.kSpring-1)
meanV['mult'].Draw("same")
meanV['2D'].SetLineColor(ROOT.kMagenta-3)
meanV['2D'].Draw('same')
leg2.Draw()


for h in fittime:
    fittime[h].GetXaxis().SetTitle('t (s)')
    fittime[h].SetStats(0)
    fittime[h].GetYaxis().SetTitle('Counts')
    fittime[h].GetYaxis().SetTitleOffset(ytitoff)
    fittime[h].GetXaxis().SetTitleOffset(xtitoff)
    fittime[h].SetLineWidth(2)
    fittime[h].GetXaxis().SetLabelOffset(labeloff)
    fittime[h].GetYaxis().SetLabelOffset(labeloff)

legt=ROOT.TLegend(0.6,0.7,0.9,0.89)
legt.SetBorderSize(0)
legt.SetFillColor(0)
legt.SetFillStyle(4000)
legt.AddEntry(fittime['mult'],'RooFit function * 100','l')
legt.AddEntry(fittime['2D'],'Free function','l')
canv.append(ROOT.TCanvas('canv3','time'))
#canv[-1].Divide(2,1)
#canv[-1].cd(1)
fittime['mult'].SetLineColor(ROOT.kSpring-1)
fittime['mult'].Draw()
fittime['2D'].SetLineColor(ROOT.kMagenta-3)
fittime['2D'].Draw('same')
legt.Draw()
#canv[-1].cd(2)

canv.append(ROOT.TCanvas('canv4','dtime'))
fittime['diff'].Draw()

for co in meanErr:
    for h in meanErr[co]:
        hist=meanErr[co][h]
        hist.SetStats(0)
        hist.GetXaxis().SetTitle('{0} (cm)'.format(co))
        hist.GetYaxis().SetTitle('Counts')
        hist.GetYaxis().SetTitleOffset(ytitoff)
        hist.GetXaxis().SetTitleOffset(xtitoff)
        hist.SetLineWidth(2)
        hist.SetTitle('')
        hist.GetXaxis().SetLabelOffset(labeloff)
        hist.GetYaxis().SetLabelOffset(labeloff)
        
leg3=ROOT.TLegend(0.6,0.7,0.9,0.89)
leg3.SetBorderSize(0)
leg3.SetFillColor(0)
leg3.SetFillStyle(4000)
leg3.AddEntry(meanErr['V']['mult'],'RooFit function','l')
leg3.AddEntry(meanErr['V']['2D'],'Free function','l')

canv.append(ROOT.TCanvas('canv5','meanErrU'))
canv[-1].cd().SetLogy()
meanErr['U']['mult'].SetLineColor(ROOT.kSpring-1)
meanErr['U']['mult'].Draw()
meanErr['U']['2D'].SetLineColor(ROOT.kMagenta-3)
meanErr['U']['2D'].Draw('same')
leg3.Draw()

canv.append(ROOT.TCanvas('canv6','meanErrV'))
canv[-1].cd().SetLogy()
meanErr['V']['mult'].SetLineColor(ROOT.kSpring-1)
meanErr['V']['mult'].Draw()
meanErr['V']['2D'].SetLineColor(ROOT.kMagenta-3)
meanErr['V']['2D'].Draw('same')
leg3.Draw()

if args.save!='Nope':
    for c in canv:
        c.SaveAs('{0}/{1}.svg'.format(args.save,c.GetTitle()))

print "Number of failed fits:",tree.GetEntries()
print "RooFit: {0} ({2:.4}%); Free fit: {1} ({3:.4}%)".format(failed['mult'],failed['2D'],float(float(failed['mult'])/tree.GetEntries())*100,float(float(failed['2D'])/tree.GetEntries())*100)
print ""
print "Average fitting time:"
print "RooFit: {0:.4}s; Free fit: {1:.4}s".format(fittime['mult'].GetMean()/100,fittime['2D'].GetMean())

u=raw_input('done?')