import ROOT, argparse, math, sys, os, copy
pandapath=os.environ.get('PANDAPATH')
sys.path.append(pandapath+'/macro/tpc/FOPI/mberger')
sys.path.append(pandapath+'/padresponse')
from padresponse import electron, electronCollection
            

parser=argparse.ArgumentParser(description='calculated the padresponse in one bin for tpc pad')
parser.add_argument('elecFile',help='the file with the electron drifts and weighting fields')
parser.add_argument('--ne',help='number of electrons to read in',type=int,default=-1)
parser.add_argument('--force',help='also the forces are inside the file, so plot them',action='store_true')
parser.add_argument('--pfile',help='the pdf file to save the canvas into',default='./padresponse_tmp.pdf')
parser.add_argument('--rfile',help='the root file to write al the hists into',type=str,default='./padresponse_tmp.root')
parser.add_argument('--batch',help='batch mode',action='store_true')
parser.add_argument('--debug',help='enable debug output',action='store_true')

args=parser.parse_args()

if args.batch:
    ROOT.gROOT.SetBatch(True)


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

canvas=ROOT.TCanvas("c","c",1000,500)
canvas.Divide(2,2)
velCanv=ROOT.TCanvas('vc','vc',1500,500)
velCanv.Divide(3,1)
forceCanvas=ROOT.TCanvas("fc",'force',1500,500)
forceCanvas.Divide(3,1)
EforceCanvas=ROOT.TCanvas("efc",'Eforce',1500,500)
EforceCanvas.Divide(3,1)
wFieldCanvas=ROOT.TCanvas("wfc",'wieghting field',1500,500)
wFieldCanvas.Divide(3,1)
posCanvas=ROOT.TCanvas("posC",'final pos',1000,500)
posCanvas.Divide(2,2)

filewords=args.elecFile.split('/')[-1]
filewords=filewords.split('_')
coll1=electronCollection(float(filewords[1]),float(filewords[2]))
coll1.readFromFile(args.elecFile,args.ne)

binwidth=coll1.electrons[0].times[-1] - coll1.electrons[0].times[-2]
print 'binwidth:',binwidth,coll1.maxt/binwidth,coll1.maxt

h=ROOT.TH1D("hc","Current",int(coll1.maxt/binwidth)+1,-binwidth,coll1.maxt+binwidth)
h.SetStats(0)
ht=ROOT.TH1D("time","time",int(coll1.maxt/binwidth)+1,-binwidth,coll1.maxt+binwidth)
ht.SetStats(0)
hwfMag=ROOT.TH2D("field","w field",int(coll1.maxt/binwidth)+1,-binwidth,coll1.maxt+binwidth,len(coll1.electrons),0,len(coll1.electrons))
hwfMag.SetStats(0)
hv=ROOT.TH2D("vel","vel",int(coll1.maxt/binwidth)+1,-binwidth,coll1.maxt+binwidth,len(coll1.electrons),0,len(coll1.electrons))
hv.SetStats(0)
hfinalpos=ROOT.TH2D("hfinalpos","Final Position",800,-0.004,0.004,800,-0.004,0.004)
hvc=[]
hf=[]
hef=[]
hwf=[]
hpos=[]
hposVsT=[]
for i in range(3):
    hvc.append(ROOT.TH2D('hv'+str(i),"vel "+str(i),int(coll1.maxt/binwidth)+1,-binwidth,coll1.maxt+binwidth,len(coll1.electrons),0,len(coll1.electrons)))
    hvc[-1].SetStats(0)
    hf.append(ROOT.TH2D('hf'+str(i),"force "+str(i),int(coll1.maxt/binwidth)+1,-binwidth,coll1.maxt+binwidth,len(coll1.electrons),0,len(coll1.electrons)))
    hf[-1].SetStats(0)
    hef.append(ROOT.TH2D('hef'+str(i),"E force "+str(i),int(coll1.maxt/binwidth)+1,-binwidth,coll1.maxt+binwidth,len(coll1.electrons),0,len(coll1.electrons)))
    hef[-1].SetStats(0)
    hwf.append(ROOT.TH2D('hwf'+str(i),"W Field "+str(i),int(coll1.maxt/binwidth)+1,-binwidth,coll1.maxt+binwidth,len(coll1.electrons),0,len(coll1.electrons)))
    hwf[-1].SetStats(0)
    hpos.append(ROOT.TH1D('hfpos'+str(i),"final position "+str(i),800,-0.004,0.004))
    hpos[-1].SetStats(0)
    hposVsT.append(ROOT.TH2D('hfposVsT'+str(i),"position vs t "+str(i),int(coll1.maxt/binwidth)+1,-binwidth,coll1.maxt+binwidth,800,-0.004,0.004))
    hposVsT[-1].SetStats(0)

#coll1.electrons[0].calcRamo()
counter=-1
for elecs in coll1.electrons:
#    elecs.calcRamo()
    counter=-1
    #print elecs.id
    hfinalpos.Fill(elecs.getPos(-1).X()-coll1.startX,elecs.getPos(-1).Y()-coll1.startY)
    for i in range(3):
        hpos[i].Fill(elecs.getPos(-1)[i]-coll1.getStart()[i])


    for val in elecs.induced:
        counter+=1
        h.Fill(elecs.times[counter],val)
        ht.Fill(elecs.times[counter])
        hwfMag.Fill(elecs.times[counter],elecs.id,elecs.getWField(counter).Mag())
        hv.Fill(elecs.times[counter],elecs.id,elecs.getVel(counter).Mag())
        for i in range(3):
            hvc[i].Fill(elecs.times[counter],elecs.id,elecs.getVel(counter)[i])
            hwf[i].Fill(elecs.times[counter],elecs.id,elecs.getWField(counter)[i])
            hposVsT[i].Fill(elecs.times[counter],elecs.getPos(counter)[i]-coll1.getStart()[i])
            if args.force:
                hf[i].Fill(elecs.times[counter],elecs.id,elecs.getForce(counter)[i])
                hef[i].Fill(elecs.times[counter],elecs.id,elecs.getEForce(counter)[i])

rfile=ROOT.TFile(args.rfile,'recreate')
canvas.cd(1)
print '1'

print '2'
h.Write()
h2=ROOT.TH1D("h2","h2",5000,0,0)
for t in coll1.current:
    h2.Fill(float(t),coll1.current[t])
h2.SetLineColor(2)
h2.Write()
h2.Draw()
#h.Draw('same')
canvas.cd(2)
ht.Draw()
ht.Write()
canvas.cd(3)
hwfMag.Draw("colz")
hwfMag.Write()
canvas.cd(4)
hv.Draw('colz')
hv.Write()
canvas.Update()
posCanvas.cd(1)
hfinalpos.Draw('colz')
hfinalpos.Write()

for i in range(3):
    velCanv.cd(i+1)
    hvc[i].Draw('colz')
    hvc[i].Write()
    wFieldCanvas.cd(i+1)
    hwf[i].Draw('colz')
    hwf[i].Write()
    posCanvas.cd(i+2)
    hpos[i].Draw()
    hpos[i].Write()
    hposVsT[i].Write()
    if args.force:
        forceCanvas.cd(i+1)
        hf[i].Draw('colz')
        hf[i].Write()
        EforceCanvas.cd(i+1)
        hef[i].Draw('colz')
        hef[i].Write()

velCanv.Update()
forceCanvas.Update()
EforceCanvas.Update()
wFieldCanvas.Update()
posCanvas.Update()

canvas.SaveAs(args.pfile+"(")
velCanv.SaveAs(args.pfile)
forceCanvas.SaveAs(args.pfile)
posCanvas.SaveAs(args.pfile)
EforceCanvas.SaveAs(args.pfile)
wFieldCanvas.SaveAs(args.pfile+")")

canvas.Write()
velCanv.Write()
posCanvas.Write()
forceCanvas.Write()
EforceCanvas.Write()
wFieldCanvas.Write()
rfile.Close()

#print "integral1:",h.Integral()
print "integral2:",coll1.getCurrentInt(True)
#print "integral3:",h2.Integral()

if not args.batch:
    u='w'
    while u!='q':
        u=raw_input('done?')