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="compare digiresiduals with theoretical values")
parser.add_argument("recofile",help='the recofile',type=str)
parser.add_argument('--outfile',help='the root output file',type=str,default='')
parser.add_argument('--events',help='number of events to process',type=int,default=-1)
parser.add_argument('--diffu',help='the logitudinal diffusion constant (sqrt(cm))',type=float,default=0.0197879)

args=parser.parse_args()


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

recofile=ROOT.TFile(args.recofile,'read')
tree=recofile.Get("cbmsim")
tree.SetBranchStatus("*",0)
tree.SetBranchStatus('CutCosmics.*',1)

ztuples=[]
zcovs=[]
zvariance=[]
zvarianceX=[]
zvarianceY=[]
znum=[]
for iz in range(75):
    ztuples.append(ROOT.TNtuple('xypos_{0}'.format(iz),'','x:y:amp'))
    #zcovs.append(ROOT.TMatrixD
    zvariance.append(0)
    znum.append(0)
    zvarianceX.append(0)
    zvarianceY.append(0)

canvas=ROOT.TCanvas('canvas','calib gem spread',1500,1000)
canvas.Divide(3,2)
hdigiresperp=ROOT.TH2D('hdigiresperp','DigiResidual Perp vs Z',75,0,75,50,-1,1)
hdigiresperpsign=ROOT.TH2D('hdigiresperpsign','DigiResidual Perp (signed) vs Z',75,0,75,50,-1,1)
hdigiresX=ROOT.TH2D('hdigiresX','DigiResidual X vs Z',75,0,75,50,-1,1)
hdigiresY=ROOT.TH2D('hdigiresY','DigiResidual Y vs Z',75,0,75,50,-1,1)
iev=-1
for e in tree:
    iev+=1
    if iev%10==0:
        print 'Event {0}'.format(iev)
    
    if args.events!=-1 and args.events<iev:
        break

    itr=-1
    for tr in e.CutCosmics:
        itr+=1
        
        for icl in range(tr.nhits()):
            
            digiTrackRes=tr.GetDigiTrackRes(icl)
            digiTrackMom=tr.GetDigiMom(icl)
            digiPos=tr.GetDigiPos(icl)

            for idi in range(len(digiTrackRes)):
                sign=1
                if (digiTrackMom[idi].DeltaPhi(digiTrackRes[idi])<0):
                    sign=-1
                hdigiresperp.Fill(digiPos[idi].Z(),digiTrackRes[idi].Perp())
                hdigiresperpsign.Fill(digiPos[idi].Z(),digiTrackRes[idi].Perp()*sign)
                hdigiresX.Fill(digiPos[idi].Z(),digiTrackRes[idi].X())
                hdigiresY.Fill(digiPos[idi].Z(),digiTrackRes[idi].Y())
                zvariance[int(digiPos[idi].Z())]+=digiTrackRes[idi].Perp()
                zvarianceX[int(digiPos[idi].Z())]+=digiTrackRes[idi].X()**2
                zvarianceY[int(digiPos[idi].Z())]+=digiTrackRes[idi].Y()**2
                znum[int(digiPos[idi].Z())]+=1


leg=ROOT.TLegend(0.1,0.8,0.4,0.9)
gvariance=ROOT.TGraph()
gvariance.SetName('graph_variance')
gvariance.SetTitle('perp deviation')
gdiffusion=ROOT.TGraph()
gdiffusion.SetName('graph_diffusion')
gdiffusion.SetTitle('diffusion')
gvarDifFrac=ROOT.TGraph()
gvarDifFrac.SetName('graph_varDifFrac')
gvarianceX=ROOT.TGraph()
gvarianceX.SetName('graph_varianceX')
gvarianceX.SetTitle('X deviation')
gvarianceY=ROOT.TGraph()
gvarianceY.SetName('graph_varianceY')
gvarianceY.SetTitle('Y deviation')

for ivar in range(len(zvariance)):
    if znum[ivar]!=0:
        zvariance[ivar]=zvariance[ivar]/znum[ivar]
        zvarianceX[ivar]=zvarianceX[ivar]/znum[ivar]
        zvarianceY[ivar]=zvarianceY[ivar]/znum[ivar]
    else:
        zvariance[ivar]=0
        zvarianceX[ivar]=0
        zvarianceY[ivar]=0
    zsig=ROOT.TMath.Sqrt(zvariance[ivar])
    zsigX=ROOT.TMath.Sqrt(zvarianceX[ivar])
    zsigY=ROOT.TMath.Sqrt(zvarianceY[ivar])
    gvarianceX.SetPoint(ivar,ivar+0.5,zsigX)
    gvarianceY.SetPoint(ivar,ivar+0.5,zsigY)
    gvariance.SetPoint(ivar,ivar+0.5,zsig)
    elecDiff=args.diffu*ROOT.TMath.Sqrt(ivar+0.5)
    gdiffusion.SetPoint(ivar,ivar+0.5,elecDiff)
    if zsig!=0:
        gvarDifFrac.SetPoint(ivar,ivar+0.5,elecDiff/zsig)

canvas.cd(1)
hdigiresperp.Draw('colz')
canvas.cd(2)
hdigiresperpsign.Draw('colz')
canvas.cd(4)
hdigiresX.Draw('colz')
canvas.cd(5)
hdigiresY.Draw('colz')
canvas.cd(3)
leg.AddEntry(gvariance,'perp deviation','p')
gvariance.SetMarkerStyle(20)
gvariance.Draw('AP')
gdiffusion.SetMarkerStyle(20)
gdiffusion.SetMarkerColor(2)
gdiffusion.Draw('P')
gvarianceX.SetMarkerStyle(20)
gvarianceX.SetMarkerColor(6)
gvarianceY.SetMarkerStyle(20)
gvarianceY.SetMarkerColor(7)
gvarianceX.Draw('P')
gvarianceY.Draw('P')
canvas.cd(6)
gvarDifFrac.SetMarkerStyle(20)
gvarDifFrac.Draw('AP')
canvas.Update()

if args.outfile!='':
    routfile=ROOT.TFile(args.outfile,'recreate')
    hdigiresperp.Write()
    hdigiresperpsign.Write()
    hdigiresX.Write()
    hdigiresY.Write()
    gvariance.Write()
    gdiffusion.Write()
    gvarDifFrac.Write()
    gvarianceX.Write()
    gvarianceY.Write()
    canvas.Write()
    routfile.Close()

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