import ROOT, glob, math, sys, os,argparse
from ROOT import std

parser=argparse.ArgumentParser(description='Analyze the Cluster pulls')
parser.add_argument('recofile',help='name of the reco file',type=str)
parser.add_argument('--AmpCut',help='Amplitude cut',type=int,default=0)
parser.add_argument('--nEvents',help='numner of events',type=int,default=0)
parser.add_argument('--SizeCut',help='',type=int,default=0)
parser.add_argument('--dres',help='enable digi residual plots',action='store_true')
parser.add_argument('--sres',help='enable sample residual plots',action='store_true')
parser.add_argument('--d',help='draw',action="store_true")
args=parser.parse_args()

recofile = args.recofile
nEvents = args.nEvents
draw = args.d
clAmpCut = args.AmpCut
clSizeCut = args.SizeCut

#argument parsing:
#for iarg in range(len(sys.argv)) :
#    arg = sys.argv[iarg]
#    if arg == "-f" :
#        recofile = sys.argv[iarg+1] 
#    if arg == "-nEvents" :
#        nEvents = int(sys.argv[iarg+1])
#    if arg == "-AmpCut" :
#        clAmpCut = int(sys.argv[iarg+1])
#    if arg == "-SizeCut" :
#        clSizeCut = int(sys.argv[iarg+1])        
#    if arg == "-d" :
#        draw = True
#finished argument parsing


ROOT.gROOT.ProcessLine(".x rootlogon.C") 
ROOT.gROOT.ProcessLine("gStyle->SetPalette(1)") 

zGem = -62.5

# Histograms
histos = []

nBins = 1000

zBins = 80
zMin = -5
zMax = 505

ClSBins = 100
ClAmpBins = 100
ClAmpMax = 100000

pullMax = 15 #10

errMax = 0.2


# Sample --------------------------------------------------------

SampleResXvsZ = ROOT.TH2D("SampleResXvsZ", "", zBins, zMin, zMax, nBins, -2,2)
histos.append(SampleResXvsZ)
SampleResYvsZ = ROOT.TH2D("SampleResYvsZ", "", zBins, zMin, zMax, nBins, -2,2)
histos.append(SampleResYvsZ)
SampleResZvsZ = ROOT.TH2D("SampleResZvsZ", "", zBins, zMin, zMax, nBins, -2,2)
histos.append(SampleResZvsZ)


# Digi --------------------------------------------------------

DigiResXvsZ = ROOT.TH2D("DigiResXvsZ", "", zBins, zMin, zMax, nBins, -2,2)
histos.append(DigiResXvsZ)
DigiResYvsZ = ROOT.TH2D("DigiResYvsZ", "", zBins, zMin, zMax, nBins, -2,2)
histos.append(DigiResYvsZ)
DigiResZvsZ = ROOT.TH2D("DigiResZvsZ", "", zBins, zMin, zMax, nBins, -2,2)
histos.append(DigiResZvsZ)

DigiResXvsclAmp = ROOT.TH2D("DigiResXvsclAmp", "", ClAmpBins, 0, ClAmpMax, nBins, -2,2)
histos.append(DigiResXvsclAmp)
DigiResYvsclAmp = ROOT.TH2D("DigiResYvsclAmp", "", ClAmpBins, 0, ClAmpMax, nBins, -2,2)
histos.append(DigiResYvsclAmp)
DigiResZvsclAmp = ROOT.TH2D("DigiResZvsclAmp", "", ClAmpBins, 0, ClAmpMax, nBins, -2,2)
histos.append(DigiResZvsclAmp)


# Cluster ------------------------------------------------------

ResXvsZ = ROOT.TH2D("ResXvsZ", "", zBins, zMin, zMax, nBins, -2,2)
histos.append(ResXvsZ)
PullXvsZ = ROOT.TH2D("PullXvsZ", "", zBins, zMin, zMax, nBins, -1.*pullMax, pullMax)
histos.append(PullXvsZ)
ErrXvsZ = ROOT.TH2D("ErrXvsZ", "", zBins, zMin, zMax, nBins, 0, errMax)
histos.append(ErrXvsZ)
ResXvsZamp = ROOT.TH2D("ResXvsZamp", "", zBins, zMin, zMax, nBins, -2,2)
histos.append(ResXvsZamp)
ResXvsZampNorm = ROOT.TH2D("ResXvsZampNorm", "", zBins, zMin, zMax, nBins, -2,2)
histos.append(ResXvsZampNorm)

ResYvsZ = ROOT.TH2D("ResYvsZ", "", zBins, zMin, zMax, nBins, -2,2)
histos.append(ResYvsZ)
PullYvsZ = ROOT.TH2D("PullYvsZ", "", zBins, zMin, zMax, nBins, -1.*pullMax, pullMax)
histos.append(PullYvsZ)

ResZvsZ = ROOT.TH2D("ResZvsZ", "", zBins, zMin, zMax, nBins, -2,2)
histos.append(ResZvsZ)
PullZvsZ = ROOT.TH2D("PullZvsZ", "", zBins, zMin, zMax, nBins, -1.*pullMax, pullMax)
histos.append(PullZvsZ)


ResXvsClSize = ROOT.TH2D("ResXvsClSize", "", ClSBins, 0, ClSBins, nBins, -2,2)
histos.append(ResXvsClSize)
PullXvsClSize = ROOT.TH2D("PullXvsClSize", "", ClSBins, 0, ClSBins, nBins, -1.*pullMax, pullMax)
histos.append(PullXvsClSize)

ResYvsClSize = ROOT.TH2D("ResYvsClSize", "", ClSBins, 0, ClSBins, nBins, -2,2)
histos.append(ResYvsClSize)
PullYvsClSize = ROOT.TH2D("PullYvsClSize", "", ClSBins, 0, ClSBins, nBins, -1.*pullMax, pullMax)
histos.append(PullYvsClSize)

ResZvsClSize = ROOT.TH2D("ResZvsClSize", "", ClSBins, 0, ClSBins, nBins, -2,2)
histos.append(ResZvsClSize)
PullZvsClSize = ROOT.TH2D("PullZvsClSize", "", ClSBins, 0, ClSBins, nBins, -1.*pullMax, pullMax)
histos.append(PullZvsClSize)


ResXvsclAmp = ROOT.TH2D("ResXvsclAmp", "", ClAmpBins, 0, ClAmpMax, nBins, -2,2)
histos.append(ResXvsclAmp)
PullXvsclAmp = ROOT.TH2D("PullXvsclAmp", "", ClAmpBins, 0, ClAmpMax, nBins, -1.*pullMax, pullMax)
histos.append(PullXvsclAmp)
ErrXvsclAmp = ROOT.TH2D("ErrXvsclAmp", "", ClAmpBins, 0, ClAmpMax, nBins, 0, errMax)
histos.append(ErrXvsclAmp)


ResYvsclAmp = ROOT.TH2D("ResYvsclAmp", "", ClSBins, 0, ClAmpMax, nBins, -2,2)
histos.append(ResYvsclAmp)
PullYvsclAmp = ROOT.TH2D("PullYvsclAmp", "", ClSBins, 0, ClAmpMax, nBins, -1.*pullMax, pullMax)
histos.append(PullYvsclAmp)

ResZvsclAmp = ROOT.TH2D("ResZvsclAmp", "", ClSBins, 0, ClAmpMax, nBins, -2,2)
histos.append(ResZvsclAmp)
PullZvsclAmp = ROOT.TH2D("PullZvsclAmp", "", ClSBins, 0, ClAmpMax, nBins, -1.*pullMax, pullMax)
histos.append(PullZvsclAmp)


ClSize2D = ROOT.TH1D("ClSize2D", "", ClSBins, 0, ClSBins)
histos.append(ClSize2D)
ClSize3D = ROOT.TH1D("ClSize3D", "", ClSBins, 0, ClSBins)
histos.append(ClSize3D)
ClAmp = ROOT.TH1D("ClAmp", "", 1000, 0, ClAmpMax)
histos.append(ClAmp)

ClSize2DvsZ = ROOT.TH2D("ClSize2DvsZ", "", zBins, zMin, zMax, ClSBins, 0, ClSBins)
histos.append(ClSize2DvsZ)
ClSize3DvsZ = ROOT.TH2D("ClSize3DvsZ", "", zBins, zMin, zMax, ClSBins, 0, ClSBins)
histos.append(ClSize3DvsZ)
ClAmpVsSize2D = ROOT.TH2D("ClAmpVsSize2D", "", ClSBins, 0, ClSBins, 1000, 0, ClAmpMax)
histos.append(ClAmpVsSize2D)
ClAmpVsSize3D = ROOT.TH2D("ClAmpVsSize3D", "", ClSBins, 0, ClSBins, 1000, 0, ClAmpMax)
histos.append(ClAmpVsSize3D)
ClAmpvsZ = ROOT.TH2D("ClAmpvsZ", "", zBins, zMin, zMax, 1000, 0, ClAmpMax)
histos.append(ClAmpvsZ)

NClvsZ = ROOT.TH1D("NClvsZ", "", zBins, zMin, zMax)
histos.append(NClvsZ)

Rfile = ROOT.TFile.Open(recofile, "read")
tree = Rfile.Get("cbmsim")

#tree.Print()

tree.SetBranchStatus("*", 0)
tree.SetBranchStatus("TpcCluster.*", 1)
tree.SetBranchStatus("MC_Sample_Residuals.*", 1)
tree.SetBranchStatus("MC_Digi_Residuals.*", 1)
tree.SetBranchStatus("MC_Cluster_Residuals.*", 1)

outname = recofile[0:recofile.rfind(".")]
outname += ".ana.root"
print outname
outfile = ROOT.TFile(outname, "recreate")



counter = 0

for e in tree :

    counter += 1;
    if (counter%100==0) :
        print "at event ", counter
        
    if (nEvents>0 and counter>nEvents):
        break


    if args.sres:
        for resCol in e.MC_Sample_Residuals :

            for iRes in range(0,resCol.getSize()) :

                res = resCol.getResidual(iRes);
                z = res.getRefPos().Z() - zGem;

                resX = res.getResXYZ().X()
                SampleResXvsZ.Fill(z, resX)

                resY = res.getResXYZ().Y()
                SampleResYvsZ.Fill(z, resY)

                resZ = res.getResXYZ().Z()
                SampleResZvsZ.Fill(z, resZ)
      
            
    if args.dres:
        for resCol in e.MC_Digi_Residuals :

            for iRes in range(0,resCol.getSize()) :

                res = resCol.getResidual(iRes);
                z = res.getRefPos().Z() - zGem;

                #digi = e.TpcDigi.At(res.getHitIndex())
                #digiAmp = cl.amp()

                resX = res.getResXYZ().X()
                DigiResXvsZ.Fill(z, resX)

                resY = res.getResXYZ().Y()
                DigiResYvsZ.Fill(z, resY)

                resZ = res.getResXYZ().Z()
                DigiResZvsZ.Fill(z, resZ)

                #DigiResXvsclAmp.Fill(digiAmp, resX)
                #DigiResYvsclAmp.Fill(digiAmp, resY)
                #DigiResZvsclAmp.Fill(digiAmp, resZ)
                        
            
    for resCol in e.MC_Cluster_Residuals :
   
        for iRes in range(0,resCol.getSize()) :
            
            res = resCol.getResidual(iRes);
            z = res.getRefPos().Z() - zGem;
            cl = e.TpcCluster.At(res.getHitIndex())
            
            clAmp = cl.amp()
            clSize2D = cl.get2DSize()
            clSize3D = cl.size();
            
            
            if (clAmp < clAmpCut or clSize2D < clSizeCut) :
                continue
            
            resX = res.getResXYZ().X()
            errX = cl.sig().X();
            ResXvsZ.Fill(z, resX)
            ResXvsZamp.Fill(z,resX,clAmp)

            PullXvsZ.Fill(z, resX/errX)
            ErrXvsZ.Fill(z, errX)
            
            resY = res.getResXYZ().Y()
            errY = cl.sig().Y();
            ResYvsZ.Fill(z, resY)
            PullYvsZ.Fill(z, resY/errY)
            
            resZ = res.getResXYZ().Z()
            errZ = cl.sig().Z();
            ResZvsZ.Fill(z, resZ)
            PullZvsZ.Fill(z, resZ/errZ)
            
            ClSize2D.Fill(clSize2D)
            ClSize3D.Fill(clSize3D)
            
            ResXvsClSize.Fill(clSize2D, resX)
            PullXvsClSize.Fill(clSize2D, resX/errX)
            
            ResYvsClSize.Fill(clSize2D, resY)
            PullYvsClSize.Fill(clSize2D, resY/errY)
            
            ResZvsClSize.Fill(clSize2D, resZ)
            PullZvsClSize.Fill(clSize2D, resZ/errZ)
            
            
            ClAmp.Fill(clAmp)
            
            ResXvsclAmp.Fill(clAmp, resX)
            PullXvsclAmp.Fill(clAmp, resX/errX)
            ErrXvsclAmp.Fill(clAmp, errX)
            
            ResYvsclAmp.Fill(clAmp, resY)
            PullYvsclAmp.Fill(clAmp, resY/errY)
            
            ResZvsclAmp.Fill(clAmp, resZ)
            PullZvsclAmp.Fill(clAmp, resZ/errZ)
            
            ClSize2DvsZ.Fill(z, clSize2D)
            ClSize3DvsZ.Fill(z, clSize3D)
            ClAmpVsSize2D.Fill(clSize2D, clAmp)
            ClAmpVsSize3D.Fill(clSize3D, clAmp)
            ClAmpvsZ.Fill(z, clAmp)
            
            NClvsZ.Fill(z)


#--------------------------

ResPullsVsZ=ROOT.TCanvas()
ResPullsVsZ.Divide(4,3)

(ResPullsVsZ.cd(1)).SetLogz()
ResXvsZ.Draw("colz")

ResPullsVsZ.cd(2)
ResXvsZ.FitSlicesY(0,0,-1,0,"QNRL")
ROOT.gDirectory.Get("ResXvsZ_2").Draw()

(ResPullsVsZ.cd(3)).SetLogz()
#PullXvsZ.Draw("colz")
ResXvsZamp.Draw("colz")

(ResPullsVsZ.cd(4)).SetLogz()
ResXvsZampNorm.Divide(ResXvsZamp,ResXvsZ)
ResXvsZampNorm.Draw("colz")
PullXvsZ.FitSlicesY(0,0,-1,0,"QNRL")
#ROOT.gDirectory.Get("PullXvsZ_2").Draw()


(ResPullsVsZ.cd(5)).SetLogz()
ResYvsZ.Draw("colz")

ResPullsVsZ.cd(6)
ResYvsZ.FitSlicesY(0,0,-1,0,"QNRL")
ROOT.gDirectory.Get("ResYvsZ_2").Draw()

(ResPullsVsZ.cd(7)).SetLogz()
PullYvsZ.Draw("colz")

ResPullsVsZ.cd(8)
PullYvsZ.FitSlicesY(0,0,-1,0,"QNRL")
ROOT.gDirectory.Get("PullYvsZ_2").Draw()


(ResPullsVsZ.cd(9)).SetLogz()
ResZvsZ.Draw("colz")

ResPullsVsZ.cd(10)
ResZvsZ.FitSlicesY(0,0,-1,0,"QNRL")
ROOT.gDirectory.Get("ResZvsZ_2").Draw()

(ResPullsVsZ.cd(11)).SetLogz()
PullZvsZ.Draw("colz")

ResPullsVsZ.cd(12)
PullZvsZ.FitSlicesY(0,0,-1,0,"QNRL")
ROOT.gDirectory.Get("PullZvsZ_2").Draw()

#--------------------------

ResPullsVs2DSize=ROOT.TCanvas()
ResPullsVs2DSize.Divide(4,3)

(ResPullsVs2DSize.cd(1)).SetLogz()
ResXvsClSize.Draw("colz")

ResPullsVs2DSize.cd(2)
ResXvsClSize.FitSlicesY(0,0,-1,0,"QNRL")
ROOT.gDirectory.Get("ResXvsClSize_2").Draw()

(ResPullsVs2DSize.cd(3)).SetLogz()
PullXvsClSize.Draw("colz")

ResPullsVs2DSize.cd(4)
PullXvsClSize.FitSlicesY(0,0,-1,0,"QNRL")
ROOT.gDirectory.Get("PullXvsClSize_2").Draw()


(ResPullsVs2DSize.cd(5)).SetLogz()
ResYvsClSize.Draw("colz")

ResPullsVs2DSize.cd(6)
ResYvsClSize.FitSlicesY(0,0,-1,0,"QNRL")
ROOT.gDirectory.Get("ResYvsClSize_2").Draw()

(ResPullsVs2DSize.cd(7)).SetLogz()
PullYvsClSize.Draw("colz")

ResPullsVs2DSize.cd(8)
PullYvsClSize.FitSlicesY(0,0,-1,0,"QNRL")
ROOT.gDirectory.Get("PullYvsClSize_2").Draw()


(ResPullsVs2DSize.cd(9)).SetLogz()
ResZvsClSize.Draw("colz")

ResPullsVs2DSize.cd(10)
ResZvsClSize.FitSlicesY(0,0,-1,0,"QNRL")
ROOT.gDirectory.Get("ResZvsClSize_2").Draw()

(ResPullsVs2DSize.cd(11)).SetLogz()
PullZvsClSize.Draw("colz")

ResPullsVs2DSize.cd(12)
PullZvsClSize.FitSlicesY(0,0,-1,0,"QNRL")
ROOT.gDirectory.Get("PullZvsClSize_2").Draw()

#--------------------------

ResPullsVsAmp=ROOT.TCanvas()
ResPullsVsAmp.Divide(4,3)

(ResPullsVsAmp.cd(1)).SetLogz()
ResXvsclAmp.Draw("colz")

ResPullsVsAmp.cd(2)
ResXvsclAmp.FitSlicesY(0,0,-1,0,"Q4NRL")
ROOT.gDirectory.Get("ResXvsclAmp_2").Draw()
histos.append(ROOT.gDirectory.Get("ResXvsclAmp_2"))

(ResPullsVsAmp.cd(3)).SetLogz()
PullXvsclAmp.Draw("colz")

ResPullsVsAmp.cd(4)
PullXvsclAmp.FitSlicesY(0,0,-1,0,"Q4NRL")
ROOT.gDirectory.Get("PullXvsclAmp_2").Draw()
histos.append(ROOT.gDirectory.Get("PullXvsclAmp_2"))

ErrXvsclAmp.FitSlicesY(0,0,-1,0,"Q4NRL")
histos.append(ROOT.gDirectory.Get("ErrXvsclAmp_0"))
histos.append(ROOT.gDirectory.Get("ErrXvsclAmp_1"))
histos.append(ROOT.gDirectory.Get("ErrXvsclAmp_2"))


(ResPullsVsAmp.cd(5)).SetLogz()
ResYvsclAmp.Draw("colz")

ResPullsVsAmp.cd(6)
ResYvsclAmp.FitSlicesY(0,0,-1,0,"Q4NRL")
ROOT.gDirectory.Get("ResYvsclAmp_2").Draw()
histos.append(ROOT.gDirectory.Get("ResYvsclAmp_2"))

(ResPullsVsAmp.cd(7)).SetLogz()
PullYvsclAmp.Draw("colz")

ResPullsVsAmp.cd(8)
PullYvsclAmp.FitSlicesY(0,0,-1,0,"Q4NRL")
ROOT.gDirectory.Get("PullYvsclAmp_2").Draw()
histos.append(ROOT.gDirectory.Get("PullYvsclAmp_2"))


(ResPullsVsAmp.cd(9)).SetLogz()
ResZvsclAmp.Draw("colz")

ResPullsVsAmp.cd(10)
ResZvsclAmp.FitSlicesY(0,0,-1,0,"Q4NRL")
ROOT.gDirectory.Get("ResZvsclAmp_2").Draw()
histos.append(ROOT.gDirectory.Get("ResZvsclAmp_2"))

(ResPullsVsAmp.cd(11)).SetLogz()
PullZvsclAmp.Draw("colz")

ResPullsVsAmp.cd(12)
PullZvsclAmp.FitSlicesY(0,0,-1,0,"Q4NRL")
ROOT.gDirectory.Get("PullZvsclAmp_2").Draw()
histos.append(ROOT.gDirectory.Get("PullZvsclAmp_2"))

#--------------------------

c4=ROOT.TCanvas()
c4.Divide(2,3)

c4.cd(1)
ClSize2D.Draw()

c4.cd(2)
ClSize3D.Draw()

c4.cd(3)
ClSize2DvsZ.Draw("colz")

c4.cd(4)
ClSize3DvsZ.Draw("colz")

c4.cd(5)
ClAmpVsSize2D.Draw("colz")

c4.cd(6)
ClAmpVsSize3D.Draw("colz")

#--------------------------

c5=ROOT.TCanvas()
c5.Divide(2,3)

c5.cd(1)
ClAmp.Draw()

c5.cd(2)
ClAmpvsZ.Draw("colz")

c5.cd(3)
NClvsZ.Draw()

#--------------------------
if draw :
    raw_input("done?")

outfile.cd()
ResPullsVsZ.Write()
ResPullsVs2DSize.Write()
ResPullsVsAmp.Write()
c4.Write()
c5.Write()
for h in histos:
    h.Write()
outfile.Close()

print "wrote out file", outname