import sys, os, copy, math
pandapath = os.environ.get('PANDAPATH')
sys.path.append(pandapath + '/macro/tpc/FOPI/python/argparse-1.2.1')
sys.path.append(pandapath + '/macro/tpc/FOPI/mberger')
import argparse, ROOT
from functions import set_palette, openTree, thisIsTheEnd, getSlices, setGraph

parser=argparse.ArgumentParser("plot driftvelocity")
parser.add_argument('infile',help='the file with the tree',type=str)
args=parser.parse_args()

rfile=ROOT.TFile(args.infile,'read')
tree=rfile.Get('gasTree')

velGraph=ROOT.TGraphErrors()
setGraph(velGraph, "", 20, 0.1, 1, "E (V/cm)", "v_Drift (cm/ns)")


vels=[]
lastfield=0
for e in tree:
    field=e.field
    vel=e.v
    vels.append(vel)
    
    if field==lastfield:
        continue
    if (field*10%10==0):
        velGraph.SetPoint(velGraph.GetN(),field+(field-lastfield)/2.,sum(vels)/len(vels))
        rms=0
        for iv in vels:
            rms+=(iv-sum(vels)/len(vels))**2
        rms/=len(vels)
        rms=math.sqrt(rms)
        velGraph.SetPointError(velGraph.GetN()-1,(field-lastfield)/2.,rms)
        vels=[]
        lastfield=field
        
    
c1=ROOT.TCanvas("cDrift",'cDrift',1400,500)
c1.Divide(2)
c1.cd(1).SetLeftMargin(0.15)
c1.cd(1).SetFillStyle(4000)
c1.cd(1)
velGraph.Draw('AP')
velGraph.GetYaxis().SetTitle("v_Drift (cm/ns)")
velGraph.GetXaxis().SetTitle("E (V/cm)")
velGraph.GetYaxis().SetTitleOffset(1.8)
velGraph.GetXaxis().SetTitleOffset(1.5)
fitfunc=ROOT.TF1("slope","pol1",350,370)
velGraph.Fit(fitfunc,"R")
v1=fitfunc.Eval(360)
v2=fitfunc.Eval(360*(1+1e-4))
print 'driftvel at 360:',v1
print 'driftvel at 360+10^-4',v2
t1=72/v1
t2=72/v2
dt=t1-t2
print "drifttime 1:",t1,"drifttime 2:",t2
print "distance diff:",t2*v1-72,dt,dt/64
field2= ((72/( (72/v1) - 64)) -fitfunc.GetParameter(0))/fitfunc.GetParameter(1)
print field2, (field2-360)/360
#pad=ROOT.TPad("pad1","",0.5,0.1,0.9,0.5)
#pad.SetFillStyle(4000)
#pad.Draw()
#pad.cd()
c1.cd(2)
velGraph2=ROOT.TGraphErrors(velGraph)
velGraph2.SetTitle("")
velGraph2.GetXaxis().SetRangeUser(360,370)
velGraph2.GetYaxis().SetRangeUser(0.002915,0.00302)
velGraph2.Draw("AP")
velGraph2.GetYaxis().SetTitle("v_Drift (cm/ns)")
velGraph2.GetXaxis().SetTitle("E (V/cm)")
velGraph2.GetYaxis().SetTitleOffset(1.9)
velGraph2.GetXaxis().SetTitleOffset(1.5)

fitfunc2=ROOT.TF1("slope2","pol1",10,30)
velGraph.Fit(fitfunc2,"R")
drifttime=2/v1
vPerp=fitfunc2.Eval(20)
perp_offset=vPerp*drifttime
print "drift distance perop to field after 2cm drift in a 20V/cm radial field: ",perp_offset,"cm"

c1.Update()
thisIsTheEnd()