import glob,math,argparse,sys
sys.path.append('$PANDAPATH/macro/tpc/FOPI/mberger')
from anaFile import anaFile
from functions import *

parser=argparse.ArgumentParser(description="plot the residual vs z graph")
parser.add_argument('anafiles',help='the list of ana files to process',type=str,nargs='+')
parser.add_argument('--hlp',help='print help',action='store_true')
parser.add_argument('--pdir',help='the picture print directory',type=str,default="not")
parser.add_argument('--MC',help='calculate with McResX',action='store_true')
parser.add_argument('--titles',help='the titles to plot in legend and canvas for the files',
                    type=str,nargs='+',default=None)
parser.add_argument('--both',help='if you want to plot Mc and fit - first gives the fit',action='store_true')
parser.add_argument('--firstbin',help='plot the first bin distibution in single canvas',action='store_true')
parser.add_argument('--prelim',help='plot preliminary into the plot',action='store_true')
parser.add_argument('--range',help='ste the tgraph range by hand',type=float,default=[-1,-1],nargs=2)
parser.add_argument('--noStats',help='dont print stats on slices',action='store_true')
parser.add_argument('--res',help='which residual direction to use?',type=str,default='X',choices=['X','Y','Z'])
parser.add_argument('--pfile',help='the pdf file to save canvas into',type=str,default='NONE')
args=parser.parse_args()

import ROOT
ROOT.gROOT.ProcessLine(".x rootlogon.C") 
#ROOT.gROOT.ProcessLine('gStyle->SetPalette(1)')
set_palette()
ROOT.gROOT.ProcessLine('gROOT->SetStyle("Plain")')

if args.hlp:
    parser.print_help()
    exit(0)


if args.titles==None:
    args.titles=[]
    for i in range(len(args.anafiles)):
        args.titles.append("no title given")

slices=[]
dslices=[]
rfiles=[]
xreshist=[]
resdata=[]
narrows=[]
broads=[]
means=[]
maxX=0
minX=999999
maxY=0
minY=999999
maxXs=[]
minXs=[]
maxYs=[]
minYs=[]
for i in range(3):
    maxXs.append(0)
    minXs.append(999999)
    maxYs.append(0)
    minYs.append(999999)

both=0  
for ifile in range(len(args.anafiles)):
    print 'doin files:',args.anafiles[ifile]
    both+=1
    rfiles.append(anaFile(args.anafiles[ifile]))
    slices.append(rfiles[-1].getZSlices())
    dslices.append(rfiles[-1].getDriftSlices())
    if args.MC:
        xreshist.append(rfiles[-1].getMcHres(args.res))
    elif args.both:
        print "both:", both
        print "Modulo:", both%2
        modulo=(both%2)
        if modulo==0:
            xreshist.append(rfiles[-1].getMcHres(args.res))
        else:
            xreshist.append(rfiles[-1].getHres(args.res))
    else:
        xreshist.append(rfiles[-1].getHres(args.res))
    if xreshist[-1]==None:
        print "no residual histo available"
        exit(0)

    narrows.append(ROOT.TGraphErrors())  
    broads.append(ROOT.TGraphErrors())  
    means.append(ROOT.TGraphErrors())  

    pcounter=0
    resdata.append([])

    for h in xreshist[ifile]:
        resdata[ifile].append(fitres(h,0))
        #print slices[-1][pcounter],resdata[-1][-1]['nsig']

        narrows[ifile].SetPoint(pcounter,dslices[-1][pcounter],resdata[-1][-1]['nsig'])
        narrows[ifile].SetPointError(pcounter,0,resdata[-1][-1]['nsigerr'])
        broads[ifile].SetPoint(pcounter,dslices[-1][pcounter],resdata[-1][-1]['bsig'])
        broads[ifile].SetPointError(pcounter,0,resdata[-1][-1]['bsigerr'])
        means[ifile].SetPoint(pcounter,dslices[-1][pcounter],resdata[-1][-1]['csig'])
        means[ifile].SetPointError(pcounter,0,resdata[-1][-1]['csigerr'])

        maxXs[0]=max(maxXs[0],resdata[-1][-1]['nsig'])
        maxXs[1]=max(maxXs[1],resdata[-1][-1]['bsig'])
        maxXs[2]=max(maxXs[2],resdata[-1][-1]['csig'])

        minXs[0]=min(minXs[0],resdata[-1][-1]['nsig'])
        minXs[1]=min(minXs[1],resdata[-1][-1]['bsig'])
        minXs[2]=min(minXs[2],resdata[-1][-1]['csig'])

        maxYs[0]=max(maxYs[0],resdata[-1][-1]['nsig'])
        maxYs[1]=max(maxYs[1],resdata[-1][-1]['bsig'])
        maxYs[2]=max(maxYs[2],resdata[-1][-1]['csig'])

        minYs[0]=min(minYs[0],resdata[-1][-1]['nsig'])
        minYs[1]=min(minYs[1],resdata[-1][-1]['bsig'])
        minYs[2]=min(minYs[2],resdata[-1][-1]['csig'])

        minX=min(minX,resdata[-1][-1]['nsig'],resdata[-1][-1]['bsig'],resdata[-1][-1]['csig'])
        maxX=max(maxX,resdata[-1][-1]['nsig'],resdata[-1][-1]['bsig'],resdata[-1][-1]['csig'])
        minY=min(minY,resdata[-1][-1]['nsig'],resdata[-1][-1]['bsig'],resdata[-1][-1]['csig'])
        maxY=max(maxY,resdata[-1][-1]['nsig'],resdata[-1][-1]['bsig'],resdata[-1][-1]['csig'])
        pcounter+=1

    narrows[-1].SetMarkerStyle(20)
    broads[-1].SetMarkerStyle(21)
    means[-1].SetMarkerStyle(22)

resgraphs=[narrows,broads,means]

maxX*=1.05
minX*=0.95
maxY*=1.05
minY*=0.95
for j in range(3):
    maxXs[j]*=1.05
    maxYs[j]*=1.05
    minXs[j]*=0.95
    minYs[j]*=0.95


for i in range(len(args.anafiles)):
#    narrows[i].GetXaxis().SetRangeUser(minXs[0],maxXs[0])
#    broads[i].GetXaxis().SetRangeUser(minXs[1],maxXs[1])
#    means[i].GetXaxis().SetRangeUser(minXs[2],maxXs[2])

    narrows[i].GetXaxis().SetRangeUser(0,75)
    broads[i].GetXaxis().SetRangeUser(0,75)
    means[i].GetXaxis().SetRangeUser(0,75)

    if args.range[0]!=args.range[1]:
        narrows[i].GetYaxis().SetRangeUser(args.range[0],args.range[1])
        broads[i].GetYaxis().SetRangeUser(args.range[0],args.range[1])
        means[i].GetYaxis().SetRangeUser(args.range[0],args.range[1])
    else:
        narrows[i].GetYaxis().SetRangeUser(minYs[0],maxYs[0])
        broads[i].GetYaxis().SetRangeUser(minYs[1],maxYs[1])
        means[i].GetYaxis().SetRangeUser(minYs[2],maxYs[2])


for ifile in range(len(args.anafiles)):
    broads[ifile].GetYaxis().SetTitle("Residual (#mum)")
    broads[ifile].GetYaxis().SetTitleOffset(1.2)
    broads[ifile].GetXaxis().SetTitle("Drift Length Z (cm)")
    
    narrows[ifile].GetYaxis().SetTitle("Residual (#mum)")
    narrows[ifile].GetYaxis().SetTitleOffset(1.2)
    narrows[ifile].GetXaxis().SetTitle("Drift Length Z (cm)")
    
    means[ifile].GetYaxis().SetTitle("Residual (#mum)")
    means[ifile].GetYaxis().SetTitleOffset(1.2)
    means[ifile].GetXaxis().SetTitle("Drift Length Z (cm)")


canvases=[]
canvases2=[]
cfirstbin=[]
legends=[]

entries=len(args.titles)+1
entrlen=0
for tit in args.titles:
    entrlen=max(entrlen,len(tit))

for i in range(3):
    canvases.append(ROOT.TCanvas())
    legends.append(ROOT.TLegend(0.1,0.9-entries*0.045,0.1+entrlen*0.015,0.9))
    legends[-1].SetFillColor(0)
plrange=0
for i in range(len(args.anafiles)):
    for h in xreshist[i]:
        plrange=max(plrange,h.GetRMS())

minmax={}
minmax['mean']={}
minmax['narrow']={}
minmax['broad']={}
for g in minmax:
    minmax[g]['min']=10000
    minmax[g]['max']=0

x,y=ROOT.Double(0),ROOT.Double(0)
for i in range(len(args.anafiles)):
    if i==4:
        narrows[i].SetMarkerColor(i+10)
        broads[i].SetMarkerColor(i+10)
        means[i].SetMarkerColor(i+10)
    else:
        narrows[i].SetMarkerColor(i+1)
        broads[i].SetMarkerColor(i+1)
        means[i].SetMarkerColor(i+1)
    if i==0:
        opt="AP"
        narrows[i].SetTitle("Narrow Gaussian")
        broads[i].SetTitle("Broad Gaussian")
        means[i].SetTitle("Mean")
    else:
        opt="P"

    if i==0:
        t=drawDiff(0,1,0)
        legends[0].AddEntry(t,'Single Electron Drift',"l")
        legends[1].AddEntry(t,'Single Electron Drift',"l")
        #legends[2].AddEntry(t,'Single Electron Drift',"l")
        t.GetXaxis().SetTitle("Drift Length (cm)")
        t.GetYaxis().SetTitle("Residual (#mu m)")
        t.GetYaxis().SetTitleOffset(1.2)

    canvases[0].SetTitle("Narrow")
    canvases[0].cd()
    legends[0].AddEntry(narrows[i],args.titles[i],"PE")
    for p in range(narrows[i].GetN()):
        narrows[i].GetPoint(p,x,y)
        minmax['narrow']['min']=min(minmax['narrow']['min'],float(y))
        minmax['narrow']['max']=max(minmax['narrow']['max'],float(y))
    narrows[0].GetYaxis().SetRangeUser(minmax['narrow']['min']*0.9,minmax['narrow']['max']*1.1)

    narrows[i].Draw(opt)
    t.Draw("same")
    if args.prelim:
        drawPrelim()

    canvases[1].SetTitle("Broad")
    canvases[1].cd()
    legends[1].AddEntry(broads[i],args.titles[i],"PE")
    for p in range(broads[i].GetN()):
        broads[i].GetPoint(p,x,y)
        minmax['broad']['min']=min(minmax['broad']['min'],float(y))
        minmax['broad']['max']=max(minmax['broad']['max'],float(y))
    broads[0].GetYaxis().SetRangeUser(minmax['broad']['min']*0.9,minmax['broad']['max']*1.1)

    broads[i].Draw(opt)
    t.Draw("same")
    if args.prelim:
        drawPrelim()

    canvases[2].SetTitle("Mean")
    canvases[2].cd()
    legends[2].AddEntry(means[i],args.titles[i],"PE")

    for p in range(means[i].GetN()):
        means[i].GetPoint(p,x,y)
        minmax['mean']['min']=min(minmax['mean']['min'],float(y))
        minmax['mean']['max']=max(minmax['mean']['max'],float(y))
    means[0].GetYaxis().SetRangeUser(minmax['mean']['min']*0.9,minmax['mean']['max']*1.1)

    means[i].Draw(opt)
    #t.Draw("same")
    if args.prelim:
        drawPrelim()

    canvases2.append(ROOT.TCanvas())
    canvases2[-1].cd()
    canvases2[-1].Divide(3,2)
    canvases2[-1].SetTitle("Slicefits_"+args.titles[i])
    counter=0            
    for h in xreshist[i]:
        canvases2[i].cd(counter+1)
        #canvases2[i].cd(counter+1).SetLeftMargin(0.15)
        if args.noStats:
            h.SetStats(ROOT.kFALSE)
        h.GetXaxis().SetRangeUser(-2.5*plrange,2.5*plrange)
        h.GetXaxis().SetTitle('Residual (#mu m)')
#        h.GetYaxis().SetTitle('Counts')
#        h.GetYaxis().SetTitleOffset(1.6)
        h.Draw()
        drawresdata(resdata[i][counter],"sigma chi/ndf S/B mean ngaus bgaus cgaus")
        drawGaussians(resdata[i][counter])
        if counter==0 and args.firstbin:
            cfirstbin.append(ROOT.TCanvas('cfirstbin'+str(i),'first_bin '+args.titles[i]))
            h.Draw()
            drawresdata(resdata[i][counter],"sigma chi/ndf S/B mean ngaus bgaus cgaus")
            drawGaussians(resdata[i][counter])
            if args.prelim:
                drawPrelim()
            cfirstbin[i].Update()
        counter+=1

    canvases2[i].Update()
    if args.pdir!='not':
        path=ROOT.TString(args.pdir+canvases2[i].GetTitle())
        path.ReplaceAll(" ","_")
        path+=".eps"
        canvases2[i].SaveAs(path.Data())
        if args.firstbin:
            path=ROOT.TString(args.pdir+cfirstbin[i].GetTitle())
            path.ReplaceAll(" ","_")
            path+=".eps"
            cfirstbin[i].SaveAs(path.Data())

for i in range(3):
    canvases[i].cd()
    legends[i].Draw()
    canvases[i].Update()
    if args.pdir!="not":
        path=ROOT.TString(args.pdir+canvases[i].GetTitle())
        path.ReplaceAll(" ","_")
        path+=".eps"
        canvases[i].SaveAs(path.Data())

if args.pfile!='NONE':
    canvases[0].SaveAs(args.pfile+'[')
    for c in canvases:
        c.SaveAs(args.pfile)
    for c in canvases2:
        c.SaveAs(args.pfile)
    canvases[0].SaveAs(args.pfile+']')

u=raw_input("quit?")
if u=='q':
    exit()