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 of an electron after drift")
parser.add_argument('digifile',help='the digifile',type=str,default="")
parser.add_argument('--nevts',help='number of events to process',type=int,default=-1)
parser.add_argument('--hlp',help='display help',action='store_true')
args=parser.parse_args()

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

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

digifile=ROOT.TFile(args.digifile,"read")
tree=digifile.Get("cbmsim")
if tree==None:
    print "no cbmsim"
    digifile.Close()
    exit()

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

nbins=150
nx=15
ny=15
hxy=ROOT.TH2D("hxy","Primary electron position xy",nbins,-nx,nx,nbins,-ny,ny)
hxydx=ROOT.TH2D("hxydx","Primary electron x drift at pos xy",nbins,-nx,nx,nbins,-ny,ny)
hxydy=ROOT.TH2D("hxydy","Primary electron y drift at pos xy",nbins,-nx,nx,nbins,-ny,ny)

hyz=ROOT.TH2D("hyz","Primary electron position yz",400,-65,15,nbins,-15,15)
hyzd=ROOT.TH2D("hyzd","Primary electron position yz",400,-65,15,nbins,-15,15)

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()
    for ielec in range(nelec):
 #       prim=e.TpcPrimaryCluster.At(iprim)
        elec=e.TpcDriftedElectron.At(ielec)
        if elec.primClustId()>=161308080:
            continue
        prim=e.TpcPrimaryCluster.At(elec.primClustId())
        dr=math.sqrt(elec.dx()*elec.dx()+elec.dy()*elec.dy())
        hxy.Fill(prim.x(),prim.y())
        hxydx.Fill(prim.x(),prim.y(),elec.dx())
        hxydy.Fill(prim.x(),prim.y(),elec.dy())
        

c1=ROOT.TCanvas("c1","e")
#c1.Divide(1,2)
hxy.Draw("Colz")
c2=ROOT.TCanvas("c2","edrift",500,1000)
c2.Divide(1,2)
hxydxnorm=Divideh(hxydx,hxy,'hxydnorm','prim elec x drift norm')
hxydynorm=Divideh(hxydy,hxy,'hxydnorm','prim elec y drift norm')
c2.cd(1)
hxydxnorm.SetStats(0)
setRange(hxydxnorm,-0.5,0.5)
hxydxnorm.Draw("colz")
c2.cd(2)
hxydynorm.SetStats(0)
setRange(hxydynorm,-0.5,0.5)
hxydynorm.Draw("colz")



u=raw_input("done?")