import ROOT, glob, math, sys, os
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
from functions import set_palette



import ROOT,argparse,math

# {{{ hist2D
class hist2D:
    def __init__(self,name,title,nx,xmin,xmax,ny,ymin,ymax):
        self.hist=ROOT.TH2D(name,title,nx,xmin,xmax,ny,ymin,ymax)
        self.hist.SetStats(0)
        self.canv=ROOT.TCanvas("c"+name,title,20,20,1000,400)
        self.size=0.07

    def setAxis(self,xtit,ytit,ztit):
        self.hist.GetXaxis().SetTitle(xtit)
        self.hist.GetXaxis().SetTitleSize(self.size)
        self.hist.GetXaxis().SetLabelSize(self.size)

        self.hist.GetYaxis().SetTitle(ytit)
        self.hist.GetYaxis().SetTitleSize(self.size)
        self.hist.GetYaxis().SetLabelSize(self.size)
        self.hist.GetYaxis().SetTitleOffset(0.6)

        self.hist.GetZaxis().SetTitle(ztit)
        self.hist.GetZaxis().SetTitleSize(self.size)
        self.hist.GetZaxis().SetLabelSize(self.size)
        self.hist.GetZaxis().SetTitleOffset(0.6)

        self.canv.SetLeftMargin(1.)
        self.canv.SetRightMargin(.15)
        self.canv.SetBottomMargin(0.15)

    def drawPrelim(self) :
        self.canv.cd()
        prelim=ROOT.TLatex()
        prelim.SetTextColor(ROOT.kGray)
        prelim.SetTextSize(0.1)
        prelim.SetTextAngle(20)
        prelim.SetNDC()
        prelim.DrawLatex(0.38,0.45,"preliminary")

    def setRange(self,mini,maxi):
        bins=(self.hist.GetNbinsX()+1)*(self.hist.GetNbinsY()+1)*(self.hist.GetNbinsZ()+1)
        for i in range(bins):
            if self.hist.GetBinContent(i)<mini:
                self.hist.SetBinContent(i,mini)
            if self.hist.GetBinContent(i)>maxi:
                self.hist.SetBinContent(i,maxi)
  
    def Draw(self,opt):
        self.canv.cd()
        self.hist.Draw("colz")

    def save(self,prefix):
        self.canv.SaveAs(prefix+self.canv.GetName()+".pdf")

# }}}
# {{{ functions
# {{{ field arrows
def getFieldArrow(r,z,field):
    point=ROOT.TVector3(r,0,z)
    fieldp=field.value2(point)
    arrow=ROOT.TArrow()
    print "stepZ:",mstepZ*stepZ,'stepR:',mstepR*stepR
   # flen=math.sqrt((mstepR*stepR)**2+(mstepZ*stepZ)**2)
    flen=min(mstepR*stepR,mstepZ*stepZ)
    arrow.SetX1(z+stepZ/2.)
    arrow.SetY1(r+stepR/2.)
    arrow.SetX2(z+stepZ/2.-(fieldp.Z()/fieldp.Mag())*flen)
    arrow.SetY2(r+stepR/2.-(fieldp.X()/fieldp.Mag())*flen)
    print fieldp.Z()/fieldp.Mag(), fieldp.X()/fieldp.Mag(), flen
    arrow.SetOption("|>")
    arrow.SetArrowSize(0.008)
    arrow.SetAngle(30)
    arrow.SetFillColor(1)
    return arrow
# }}}
# {{{ field line
def getFieldlineRZ(r0,z0,field):
    fieldline=ROOT.TPolyLine()
    r=r0
    z=z0
    pcount=0
    fieldline.SetPoint(pcount,z,r)
#    print "new fieldline"
#    while (r<maxR and z<maxZ and r>minR and z>minZ):
    while pcount<10*nstepZ:
        oldr=r
        oldz=z

        fieldline.SetPoint(pcount,z,r)
        point=ROOT.TVector3(r,0,z)
        fieldp=field.value2(point)
        if fieldp[0]==0 and fieldp[2]==0:
            break

        r-=fieldp[0]
        z-=fieldp[2]
        
        fr=fieldp[0]
        fz=fieldp[2]
        while abs(oldr-r)>stepR or abs(oldz-z)>stepZ:
#            print '%(1)1.10f  %(2)1.10f ' % {"1":abs(oldr-r) ,"2": abs(oldz-z)}
            fr/=2
            r+=fr
            fz/=2
            z+=fz

        pcount+=1
#        print "or:",oldr,"oz:",oldz,"r:",r,"z:",z
       
        if (r>maxR or z>maxZ or r<minR or z<minZ):
            break
    return fieldline
# }}}
# {{{ my_range
def my_range(start,end,step):
    while start<=end:
        yield start
        start+=step
# }}}
# }}}


parser=argparse.ArgumentParser(description='Plot the electric field given in a file (TpcEField encoding)')
parser.add_argument('filename',help='the file with the electric field')
parser.add_argument('--nr',help='bins in r',type=int,default=0)
parser.add_argument('--nz',help='bins in z',type=int,default=0)
parser.add_argument('--rangeR',help='range in r',type=float,nargs=2,default=[-1,-1])
parser.add_argument('--rangeZ',help='range in z',type=float,nargs=2,default=[-1,-1])
parser.add_argument('--norange',help='do not rerang plots',action="store_true")
parser.add_argument('--fl',help='plot fieldlines',action="store_true")
parser.add_argument('--arr',help='plot field arrows',action='store_true')
parser.add_argument('--out',help='plot field to eps files',action="store_true")
parser.add_argument('--rout',help='root file to write the histos into',type=str,default='')
parser.add_argument('--zfieldRange',help='field plot range',type=float, nargs=2, default=[290,320])

args=parser.parse_args()

ROOT.gROOT.SetStyle("Plain")
ROOT.gROOT.ProcessLine(".x rootlogon.C") 


set_palette()
efile=open(args.filename,'r')
field=ROOT.TpcEFieldCyl(args.filename)

line=efile.readline()
words=line.split()
nomEField=float(words[0])
minR=float(words[1])
stepR=float(words[2])
nstepR=int(words[3])
minZ=float(words[4])
stepZ=float(words[5])
nstepZ=int(words[6])

maxR=minR+stepR*(nstepR-1)
maxZ=minZ+stepZ*(nstepZ-1)

print "original field has",nstepR," steps in R. starting from:",minR,"to",maxR,"in",stepR,"wide steps in R"
print "original field has",nstepZ," steps in Z. starting from:",minZ,"to",maxZ,"in",stepZ,"wide steps in Z"

if args.nr!=0:
    nstepR=args.nr
    stepR=(maxR-minR)/nstepR
    print 'new steps in R:',nstepR,' new stepsize:',stepR
if args.nz!=0:
    nstepZ=args.nz
    stepZ=(maxZ-minZ)/nstepZ

if args.rangeR[0]!=-1:
    minR=args.rangeR[0]
    maxR=args.rangeR[1]
    stepR=(maxR-minR)/nstepR

if args.rangeZ[0]!=-1:
    minZ=args.rangeZ[0]
    maxZ=args.rangeZ[1]
    stepZ=(maxZ-minZ)/nstepZ

print "plotted field has",nstepR," steps in R. starting from:",minR,"to",maxR,"in",stepR,"wide steps in R"
print "plotted field has",nstepZ," steps in Z. starting from:",minZ,"to",maxZ,"in",stepZ,"wide steps in Z"

hists=[]
hFieldR=hist2D("hFieldR","E-Field R-Component",nstepZ,minZ,maxZ,nstepR,minR,maxR)
hists.append(hFieldR)
hFieldZ=hist2D("hFieldZ","E-Field Z-Component",nstepZ,minZ,maxZ,nstepR,minR,maxR)
hists.append(hFieldZ)
hFieldAbs=hist2D("hFieldAbs","Absolute Field",nstepZ,minZ,maxZ,nstepR,minR,maxR)
hists.append(hFieldAbs)

maxfieldZ=0
minfieldZ=999999
maxfieldR=0
minfieldR=999999
for nr in range(nstepR):
    for nz in range(nstepZ):
        if (minZ+(nz+0.5)*stepZ)>maxZ:
            continue
        pos=ROOT.TVector3(float(minR+(nr+0.5)*stepR),0.,float(minZ+(nz+0.5)*stepZ))
        value=ROOT.TVector3()
        value=field.value2(pos)
        hFieldR.hist.SetBinContent(nz+1,nr+1,value.X())
        hFieldZ.hist.SetBinContent(nz+1,nr+1,value.Z())
        hFieldAbs.hist.SetBinContent(nz+1,nr+1,value.Mag())
        maxfieldR=max(value.X(),maxfieldR)
        minfieldR=min(value.X(),minfieldR)
        maxfieldZ=max(value.Z(),maxfieldZ)
        minfieldZ=min(value.Z(),minfieldZ)

hFieldR.setAxis("z (cm)","r (cm)","Er (V/cm)")
hFieldR.setRange(-30,30)
hFieldR.Draw("colz")
hFieldZ.setAxis("z (cm)","r (cm)","Ez (V/cm)")
#hFieldZ.setRange(340,380)
if not args.norange:
    hFieldZ.setRange(args.zfieldRange[0],args.zfieldRange[1])
hFieldZ.Draw("colz")
hFieldAbs.setAxis("z (cm)","r (cm)","E (V/cm)")
#hFieldAbs.setRange(340,380)
if not args.norange:
    hFieldAbs.setRange(args.zfieldRange[0],args.zfieldRange[1])
hFieldAbs.Draw("colz")
hFieldAbs.canv.cd()
lines=[]
if args.fl:
    for i in my_range(minR,maxR+stepR,stepR):
        lines.append(getFieldlineRZ(i,maxZ,field))
#        lines[-1].Print()
        lines[-1].Draw()
        print "plotting fieldline ",i,"to go:",(maxR+stepR-i)/stepR
        hFieldAbs.canv.Update()

arrows=[]
mstepR=5
mstepZ=20
if args.arr:
    for i in my_range(minR,maxR,mstepR*stepR):
        for j in my_range(minZ,maxZ,mstepZ*stepZ):
            arrows.append(getFieldArrow(i,j,field))
            arrows[-1].Draw()
            hFieldAbs.canv.Update()

if args.out:
    for h in hists:
        h.save(args.filename.replace('.txt',''))

if args.rout!='':
    rfile=ROOT.TFile(args.rout,'RECREATE')
    for h in hists:
        h.hist.Write()
        
u=raw_input("done?")