import ROOT,math,argparse,os,math
ROOT.gROOT.ProcessLine(".x rootlogon.C") 

parser=argparse.ArgumentParser(description="plot the drift path of an electron")
parser.add_argument('filename',help='file with the coordinates')
parser.add_argument('fieldfilename',help='file with the electric field')
parser.add_argument('--mean',help="plot the mean",action="store_true")
parser.add_argument('--mean2',help="plot the mean",action="store_true")
parser.add_argument('--poswise',action="store_true")
parser.add_argument('--setmax',action="store_true")
parser.add_argument('--ev',type=int,default=-1)
args=parser.parse_args()

dfile=open(args.filename,'r')
field=ROOT.TpcEFieldCyl(args.fieldfilename)
zbin=0
rbin=0

x=0
y=0
z=0
xbins=106*20
zbins=723*10
#hrz=ROOT.TH2D("hrz","rz position",zbins,17.249,17.251,zbins,5.048,5.056)
#hrz=ROOT.TH2D("hrz","rz position",zbins,0,72.3,zbins,5,15.6)
#hxy=ROOT.TH2D("hxy","xy position",xbins,5.0494,5.054,xbins,0,.0022)
#hxy=ROOT.TH2D("hxy","xy position",xbins,5.0,15.6,xbins,0,.003)
grz=ROOT.TGraph()
gxy=ROOT.TGraph()
grzm=ROOT.TGraph()
gxym=ROOT.TGraph()
dummy1=ROOT.TGraph()
dummy2=ROOT.TGraph()
hfr=ROOT.TH2D()
hfz=ROOT.TH2D()
hfabs=ROOT.TH2D()
#pos=ROOT.TVector3()

mean=[]
pos=[]
for i in range(4):
    mean.append([])
    pos.append([])

c=ROOT.TCanvas("c","bla",600,1200)
c.Divide(1,3)
c2=ROOT.TCanvas("c2","bla2",600,400)
counter=0
poscounter=-1

    
def FillGraphs():
    grz.Set(0)
    gxy.Set(0)
    grzm.Set(0)
    gxym.Set(0)

    binsz=500
    binsr=500
    minz=min(pos[2]+mean[2])
    maxz=max(pos[2]+mean[2])
    minr=min(pos[3]+mean[3])
    maxr=max(pos[3]+mean[3])
    stepr=(maxr-minr)/float(binsr)
    stepz=(maxz-minz)/float(binsz)
    hfr=ROOT.TH2D("hfieldr","Field r-component",binsr,minz,maxz,binsz,minr,maxr)
    hfz=ROOT.TH2D("hfieldz","Field z-component",binsr,minz,maxz,binsz,minr,maxr)
    hfabs=ROOT.TH2D("hfielda","Absolute Field",binsr,minz,maxz,binsz,minr,maxr)
    hfr.SetStats(0)
    hfz.SetStats(0)
    hfabs.SetStats(0)

#    grz.SetPoint(0,min(pos[2]+mean[2]),min(pos[3]+mean[3]))
#    grz.SetPoint(1,max(pos[2]+mean[2]),max(pos[3]+mean[3]))
    gxy.SetPoint(0,min(pos[0]+mean[0]),min(pos[1]+mean[1]))
    gxy.SetPoint(1,max(pos[0]+mean[0]),max(pos[1]+mean[1]))
    for r in range (binsr):
        for z in range(binsz):
            posv=ROOT.TVector3(float(minr+(r+0.5)*stepr),0.,float(minz+(z+0.5)*stepz))
            hfr.SetBinContent(z+1,r+1,field.value2(posv).X())
            hfz.SetBinContent(z+1,r+1,field.value2(posv).Z())
            hfabs.SetBinContent(z+1,r+1,field.value2(posv).Mag())
    for i in range(len(pos[0])):
        grz.SetPoint(i+2,pos[2][i],pos[3][i])
        gxy.SetPoint(i+2,pos[0][i],pos[1][i])
        grzm.SetPoint(i+2,mean[2][i],mean[3][i])
        gxym.SetPoint(i+2,mean[0][i],mean[1][i])
        if args.poswise and i%10==0:
            DrawDaShit()
    return hfr, hfz, hfabs

def DrawDaShit():
    c.cd(1)
    if not args.mean2:
        hfr.Draw("colz")
        grz.Draw("P")
        c.cd(2)
        hfz.Draw("colz")
        grz.Draw("P")
    if args.mean or args.mean2:
        grzm.SetLineColor(ROOT.kRed)
        grzm.SetMarkerColor(ROOT.kRed)
        c.cd(1)
        grzm.Draw("P")
        c.cd(2)
        grzm.Draw("P")
    c.cd(3)
    if not args.mean2:
        gxy.Draw("AP")
    if args.mean or args.mean2:
        gxym.SetLineColor(ROOT.kRed)
        gxym.SetMarkerColor(ROOT.kRed)
        gxym.Draw("P")
    c.Update()
  
    c2.cd()
    hfabs.Draw("colz")
    grz.Draw("P")

for line in dfile:
    words=line.split(";")
    if words[0]=='position':
        poscounter+=1

        mean[0].append(float(words[4])*100)
        mean[1].append(float(words[5])*100)
        mean[2].append(float(words[6])*100)
        mean[3].append(float(words[7])*100)
            
        if not args.mean2:
            pos[0].append(float(words[1])*100)
            pos[1].append(float(words[2])*100)
            pos[2].append(float(words[3])*100)
            pos[3].append(math.sqrt(pos[0][-1]*pos[0][-1]+pos[1][-1]*pos[1][-1]))

    if line.find('rPos')!=-1 and poscounter!=-1:
        poscounter=-1
        counter+=1
        if counter!=args.ev and args.ev!=-1:
            continue
        print "waiting for update"
        hfr,hfz,hfabs=FillGraphs()
        DrawDaShit()

        u=raw_input("updated "+str(counter))
        for i in range(4):
            del pos[i][:]
            del mean[i][:]
 #   else:
 #       for i in range(4):
 #           del pos[i][:]
 #           del mean[i][:]
 #       counter+=1

hfr,hfz,hfabs=FillGraphs()
DrawDaShit()
u=raw_input("updated "+str(counter+1))