import sys
sys.path.append('/home/mberger/fopiroot/fopiroot_dev/macro/tpc/FOPI/mberger')
import ROOT,glob,math,argparse
from functions import Divideh,set_palette,setRange
parser=argparse.ArgumentParser(description="plos deviation from MC")
parser.add_argument('mcfile',help='the mcfile',type=str,default="")
parser.add_argument('digifile',help='the digifile',type=str,default="")
parser.add_argument('recofile',help='the recofile',type=str,default="")
#parser.add_argument('--ampweight',help='weight with amplitude',action='store_true')
parser.add_argument('--nevts',help='number of events to process',type=int,default=-1)
parser.add_argument('--hlp',help='display help',action='store_true')
parser.add_argument('--outfile',help='the output root file with all the histos',type=str,default='')
args=parser.parse_args()

set_palette()
ROOT.gROOT.ProcessLine(".x rootlogon.C") 
ROOT.gROOT.ProcessLine('gROOT->SetStyle("Plain")')


recofile=ROOT.TFile(args.recofile,'read')
tree=recofile.Get("cbmsim")
tree.AddFriend("cbmsim",args.mcfile)
tree.AddFriend("cbmsim",args.digifile)

tree.SetBranchStatus("*",0)
tree.SetBranchStatus("TpcPrimaryCluster.*",1)
tree.SetBranchStatus("TpcDriftedElectron.*",1)
tree.SetBranchStatus('CutCosmics.*',1)
tree.SetBranchStatus('TpcPoint.*',1)

canvas=ROOT.TCanvas('canvas','MC Res',2000,1000)
canvas.Divide(4,2)
canvas2=ROOT.TCanvas('canvas2','Digi Len',1000,1000) 
canvas2.Divide(2,2)

stages=['prim','delec','digi','digiAmp']
zmax=[75,32000,75,75]
zbins=[75,80,75,75]

hresx={}
hresy={}
hresz={}
counter=-1
for stage in stages:
    counter+=1
    hresx[stage]=ROOT.TH2D('hresx_'+stage,'Residual in X '+stage,zbins[counter]+1,-1,zmax[counter],100,-1,1)
    hresy[stage]=ROOT.TH2D('hresy_'+stage,'Residual in Y '+stage,zbins[counter]+1,-1,zmax[counter],100,-1,1)
    hresz[stage]=ROOT.TH2D('hresz_'+stage,'Residual in z '+stage,zbins[counter]+1,-1,zmax[counter],100,-1,1)
    print 'generated hist for',stage
    
hdigiResZvsLength=ROOT.TH2D("hdigiResZVsLength",'digi residual Z vs length',100,0,100,100,-1,1)
hdigiLegthVsZ=ROOT.TH2D('hdigilengthsVsZ','Digi Length vs Z',75,0,75,30,0,30)
hresXY=[]
hresXYAmp=[]
for iz in range(75):
    hresXY.append(ROOT.TH2D('hresxy_{0}'.format(iz),'Residual in XY z-bin {0}'.format(iz),50,-1,1,50,-1,1))
    hresXYAmp.append(ROOT.TH2D('hresxyAmp_{0}'.format(iz),'Residual in XY Amp z-bin {0}'.format(iz),50,-1,1,50,-1,1))
ev=-1
for e in tree:
    ev+=1
    if args.nevts!=-1 and ev>args.nevts:
        break
    nprim=e.TpcPrimaryCluster.GetEntries()
    nelec=e.TpcDriftedElectron.GetEntries()
    print 'at Event:',ev

    for ielec in range(nelec):
        elec=e.TpcDriftedElectron.At(ielec)
        if elec.primClustId()>=161308080:
            continue
        prim=e.TpcPrimaryCluster.At(elec.primClustId())
        mchit=e.TpcPoint.At(prim.mcHitId())
        hresx['prim'].Fill(prim.z(),mchit.GetX()-prim.x())
        hresy['prim'].Fill(prim.z(),mchit.GetY()-prim.y())
        hresz['prim'].Fill(prim.z(),mchit.GetZ()-prim.z())

        hresx['delec'].Fill(elec.t(),mchit.GetX()-elec.x())
        hresy['delec'].Fill(elec.t(),mchit.GetY()-elec.y())
        #hresz['delec'].Fill(elec.t(),mchit.GetZ()-elec.t()) makes no sense like this

    for tr in e.CutCosmics:
        for icl in range(tr.nhits()):
            digiMcRes=tr.GetDigiMCRes(icl)
            digiPos=tr.GetDigiPos(icl)
            digiAmp=tr.GetDigisOnPlaneAmp(icl)
            digilen=tr.GetDigiLen(icl)
            dcounter=-1
            for mcres in digiMcRes:
                dcounter+=1
                #hresx['digi'].Fill(digiPos[counter].Z(),ROOT.gRandom.Gaus(mcres.X(),0.1299))
                #hresy['digi'].Fill(digiPos[counter].Z(),ROOT.gRandom.Gaus(mcres.Y(),0.1125))
                hresx['digi'].Fill(digiPos[counter].Z(),mcres.X()+ROOT.gRandom.Uniform(-0.1299,0.1299))
                hresy['digi'].Fill(digiPos[counter].Z(),mcres.Y()+ROOT.gRandom.Uniform(-0.1125,0.1125))
                hresz['digi'].Fill(digiPos[counter].Z(),mcres.Z()+ROOT.gRandom.Uniform(-0.152867/2.,0.152867/2.))
                #hresx['digiAmp'].Fill(digiPos[counter].Z(),ROOT.gRandom.Gaus(mcres.X(),0.1299),abs(digiAmp[counter]))
                #hresy['digiAmp'].Fill(digiPos[counter].Z(),ROOT.gRandom.Gaus(mcres.Y(),0.1125),abs(digiAmp[counter]))
                hresx['digiAmp'].Fill(digiPos[counter].Z(),mcres.X()+ROOT.gRandom.Uniform(-0.1299,0.1299),abs(digiAmp[counter]))
                hresy['digiAmp'].Fill(digiPos[counter].Z(),mcres.Y()+ROOT.gRandom.Uniform(-0.1125,0.1125),abs(digiAmp[counter]))
                hresz['digiAmp'].Fill(digiPos[counter].Z(),mcres.Z()+ROOT.gRandom.Uniform(-0.152867/2.0,0.152867/2.0),abs(digiAmp[counter]))
                hresXY[int(digiPos[counter].Z())].Fill(mcres.X()+ROOT.gRandom.Uniform(-0.1299,0.1299),mcres.Y()+ROOT.gRandom.Uniform(-0.1125,0.1125))
                hresXYAmp[int(digiPos[counter].Z())].Fill(mcres.X()+ROOT.gRandom.Uniform(-0.1299,0.1299),mcres.Y()+ROOT.gRandom.Uniform(-0.1125,0.1125),abs(digiAmp[counter]))

                hdigiResZvsLength.Fill(digilen[counter],mcres.Z()+ROOT.gRandom.Uniform(-0.152867/2.0,0.152867/2.0))
                hdigiLegthVsZ.Fill(digiPos[counter].Z(),digilen[counter])
                       
counter=-1
for stage in stages:
    counter+=1
    canvas.cd(counter+1)
    hresx[stage].Draw('colz')
    canvas.cd(counter+5)
    hresy[stage].Draw('colz')
canvas.Update()

canvas2.cd(1)
hdigiResZvsLength.Draw('colz')
canvas2.cd(2)
hdigiLegthVsZ.Draw('colz')
canvas2.cd(3)
hdigiLegthVsZ.FitSlicesY()
graph=ROOT.gDirectory.Get("hdigilengthsVsZ_1")
graph.Draw()
canvas2.cd(4)
graph2=hdigiLegthVsZ.ProfileX()
graph2.Draw()
canvas2.Update()

if args.outfile!='':
    routfile=ROOT.TFile(args.outfile,'recreate')
    canvas.Write()
    for stage in stages:
        hresx[stage].Write()
        hresy[stage].Write()
    routfile.mkdir('ResXY')
    routfile.cd('ResXY')
    for h in hresXY:
        h.Write()
    routfile.mkdir('ResXY_Amp')
    routfile.cd('ResXY_Amp')
    for h in hresXYAmp:
        h.Write()

    routfile.Close()
u='w'
while u!='q':
    u=raw_input('done?')