import ROOT, argparse, math, sys, os

def calcMCres(pos,e,dperc):
    mindist=999
    zdist=99
    minpoint=None
    count=-1
    npoints=e.TpcPoint.GetEntriesFast()
    for point in e.TpcPoint:
        count+=1
        if int((float(count)/float(npoints) )*100)%25==0:
            print str(dperc)+"%",str(int((float(count)/float(npoints))*100))+"% \r",

#        if point.GetZ()>1.1*pos.Z() or point.GetZ()<0.9*pos.Z():
#            continue
        pPos=ROOT.TVector3(point.GetX(),point.GetY(),point.GetZ())
        dist=(pPos-pos)
        if dist.Mag()<mindist:
            mindist=dist.Mag()
            minpoint=point
            zdist=dist.Z()
    count+=1
    pPos=ROOT.TVector3(minpoint.GetX(),minpoint.GetY(),minpoint.GetZ())
    pDir=minpoint.GetMomentum()
    A=(pDir * (pos - pPos))
    B=pDir
    B*=A
    POCA = pPos + B

    POCA_t=(((POCA.Z()-args.fzGem)/args.vd)-args.ft0)/args.ftbin

    print str(dperc)+"%",str(int((float(count)/float(npoints))*100))+"% \r",
    return (pos-POCA).Z(),POCA_t


parser=argparse.ArgumentParser(description='calibrate zoffset in PSA')
parser.add_argument('mcfile',help='path to mc file')
parser.add_argument('digifile',help='path to digi file')
parser.add_argument('--recofile',help='path to the recofile. enables cluster residual')
parser.add_argument('--outfile',help='name of the ouputfile',type=str,default="zcalib.root")
parser.add_argument('--events',help='how many events to process',type=int,default=-1)
parser.add_argument('--ft0',help='ft for digi z calc',type=float,default=0)
parser.add_argument('--ftbin',help='ftbin for digi z calc',type=float,default=64.3087)
parser.add_argument('--vd',help='drift velocity for digi z calc',type=float,default=0.00237708)
parser.add_argument('--fzGem',help='z gem for digi z calc',type=float,default=0)
parser.add_argument('--padf',help='the padplane file to use',type=str,default='tpc/parfiles/padPlane_FOPI.dat')

args=parser.parse_args()

ROOT.gROOT.ProcessLine(".x rootlogon.C") 
ROOT.gROOT.ProcessLine("gStyle->SetPalette(1)") 
ROOT.gROOT.LoadMacro("stlPYROOT.h+")

c1=ROOT.TCanvas("c1","Zoff Calib",1000,1000)
c1.Divide(2,2)

DigiMapper = []
file = open(args.padf, "r")
lines=file.readlines()
file.close()
lines.reverse()
for line in lines:
    koord = line.split(" ") 
    DigiMapper.append([koord[3], koord[4]])

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

tree.SetBranchStatus("TpcDigi.*",1)
tree.SetBranchStatus("TpcPoint.*",1)

hzoff=ROOT.TH1D("hzoff","digi-mc Residual Z",200,-10,10)
hzoffVsZ=ROOT.TH2D("hzoffVsZ","digi-mc Residual Z vs Z",75,0,75,200,-5,5)
htoffVst=ROOT.TH2D("htoffVst","digi-mc Residual t vs t",2048,0,4096,200,-100,100)

hzoffC=ROOT.TH1D("hzoffC","cluster-mc Residual Z",200,-1,1)
hzoffVsZC=ROOT.TH2D("hzoffVsZC","cluster-mc Residual Z vs Z",75,0,75,200,-2,2)

ev=-1
for e in tree:
    ev+=1
    ndigis=e.TpcDigi.GetEntriesFast()
    print "\nEvent:",ev,"(",ndigis,",",e.TpcPoint.GetEntriesFast(),")"
    if args.events!=-1 and ev>=args.events:
        break
    dcount=0
    if ndigis>1200 or e.TpcPoint.GetEntriesFast()>1500:
        print "Too Much points. skipping"
        continue

    for digi in e.TpcDigi:
        dcount+=1
        dperc=int((float(dcount)/float(ndigis))*100)

        dx=float(DigiMapper[digi.padId()][0])
        dy=float(DigiMapper[digi.padId()][1])
        dz=(digi.t()*args.ftbin+args.ft0)*args.vd+args.fzGem
        ResZ,POCA_T=calcMCres(ROOT.TVector3(dx,dy,dz),e,dperc)
        hzoff.Fill(ResZ)
        hzoffVsZ.Fill(dz,ResZ)
        htoffVst.Fill(digi.t(),POCA_T-digi.t())

    #for tr in e.CutCosmics
        

c1.cd(1)
htoffVst.Draw("colz")
c1.cd(2)
hzoffVsZ.SetStats(0)
hzoffVsZ.Draw("colz")
c1.cd(4)
profxz=hzoffVsZ.ProfileX()
profxz.Draw()
c1.cd(3)
profxt=htoffVst.ProfileX()
profxt.Draw()
c1.Update()
outfile=ROOT.TFile(args.outfile,"recreate")
c1.Write()
profxt.Write()
profxz.Write()
hzoff.Write()
htoffVst.Write()
hzoffVsZ.Write()
outfile.Close()
u='w'
while u!='q':
    u=raw_input("done?")