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

import argparse, ROOT, math
from operator import itemgetter


def setCont(h,nlevels,maxi,mini):
    h.SetContour(nlevels)
    for i in range(1,nlevels):
       h.SetContourLevel(i,mini + ((maxi-mini)/(nlevels-1))*i)
       print i ,mini + ((maxi-mini)/(nlevels-1))*i
    h.SetContourLevel(0,0.1)
    return

def setRange(h,mini,maxi):
    bins=(h.GetNbinsX()+1)*(h.GetNbinsY()+1)*(h.GetNbinsZ()+1)
    for i in range(bins):
        if h.GetBinContent(i)<mini:
            h.SetBinContent(i,mini)
        if h.GetBinContent(i)>maxi:
            h.SetBinContent(i,maxi)
    return

def setAxis(h,x,y,z=""):
    size=0.05
    h.GetXaxis().SetTitle(x)
    h.GetXaxis().SetTitleSize(size)
    h.GetXaxis().SetLabelSize(size)

    h.GetYaxis().SetTitle(y)
    h.GetYaxis().SetTitleSize(size)
    h.GetYaxis().SetLabelSize(size)
    h.GetYaxis().SetTitleOffset(0.6)

    h.GetZaxis().SetTitle(z)
    h.GetZaxis().SetTitleSize(size)
    h.GetZaxis().SetTitleOffset(0.6)
    h.GetZaxis().SetLabelSize(size)
    return

def setMargin(c):
#    c.SetLeftMargin(1.)
#    c.SetBottomMargin(0.15)
    c.cd(1).SetRightMargin(0.15)
    return

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

parser=argparse.ArgumentParser(description='Reads a COMSOL field file and creates another field file.')
parser.add_argument('filename',help='name of the COMSOL field file')
parser.add_argument('--out',help='output filename',default='')
parser.add_argument('--pr',help='save canvas to pdf',action='store_true')
parser.add_argument('--prpath',help='path to save the pictures to',type=str,default='./')
parser.add_argument('--prel',help='print preliminary on all plots',action='store_true')
parser.add_argument('--batch',help='batch mode',action='store_true')
parser.add_argument("--overwrite",help="overwrite old field file",action='store_true')
parser.add_argument('--ntit',help='no histogram titles',action='store_true')
parser.add_argument('--zfieldRange',help='field plot range',type=float, nargs=2, default=[300,500])
parser.add_argument('--smooth0',help='smooth away field values below 0',action='store_true')

args = parser.parse_args()

set_palette()
ROOT.gROOT.ProcessLine(".x rootlogon.C") 
if args.batch:
    ROOT.gROOT.SetBatch(True)


if args.filename.find('out.txt')!=-1:
    print args.filename,"not a comsol file, exiting"
    exit()

if not args.overwrite:
    if os.path.isfile(args.filename.replace('.txt','_out.txt')):
        print args.filename,"already done, exiting"
        exit()

file=open(args.filename,'r')

field=[]

rmin=1000
rmax=0
zmin=1000
zmax=0
nr=0
nz=0
fieldz=[]


oldr=-1
oldz=-1
first=True
ir=0
iz=0

print "doing file",args.filename
for line in file:
    words=line.split()
    if words[0]=='%':
        continue
    r=round(float(words[0])/10.,2)
    z=round(float(words[1])/10.,2)
    
    if first:
        oldr=r
        oldz=z
        first=False
        
    if oldr==r:
        nz+=1
    if oldz==z:
        nr+=1
        
    rmin=min(rmin,r)
    rmax=max(rmax,r)
    zmin=min(zmin,z)
    zmax=max(zmax,z)
#    0=r[mm];1=z[mm];2=Er[V/cm],3=Ez[V/cm];4=V[V]
    field.append([r,z,float(words[2]),float(words[3]),float(words[4])])
    
    #if r!=oldr:
    #    fieldz.append([])
    #    oldr=r
        
    #fieldz[ir].append([])

print "nr=",nr,"nz=",nz

if args.smooth0:
    for i in range(len(field)):
        #if 1:
        if field[i][3]<1e-10:
            rpos=field[i][0]
            zpos=field[i][1]
            ir=i%nr
            iz=int(i/nr)
            #print "ir={0}, iz={1}, rpos={2}, zpos={3}".format(ir,iz,rpos,zpos) 
            neighboars=[]
            weights=[]
            if ir<nr:
                neighboars.append(i+1)
                weights.append(1)
                if iz>0:
                    neighboars.append(i+1-nr)
                    weights.append(1/math.sqrt(2))
                if iz<nz-1:
                    neighboars.append(i+1+nr)
                    weights.append(1/math.sqrt(2))
            if ir>0:
                neighboars.append(i-1)
                weights.append(1)
                if iz>0:
                    neighboars.append(i+-1-nr)
                    weights.append(1/math.sqrt(2))
                if iz<nz-1:
                    neighboars.append(i-1+nr)
                    weights.append(1/math.sqrt(2))    
                                
            if iz>0:
                neighboars.append(i-nr)
                weights.append(1)
            if iz<nz-1:
                neighboars.append(i+nr)
                weights.append(1)
                
            newfield=0
            weight=0
            n=-1
            for nei in neighboars:
                n+=1
                if field[nei][3]>1e-10:
                    newfield+=field[nei][3]*weights[n]
                    weight+=weights[n]
            field[i][3]=newfield/weight
 
print "done reading the files",len(field)
print field[0][0],field[1][0]
path=args.filename.split('/')
filewords=path[-1].split('_')
dr=abs(field[0][0]-field[1][0])
dz=dr#abs(z[0]-z[1])

nomEz=abs(float(filewords[3]))

print rmax,rmin,dr,abs(rmax-rmin)
stepsr=int(abs(rmax-rmin)/dr)+1
print "steps in R:",stepsr
stepsz=int(abs(zmax-zmin)/dz)+1
print "steps in z:",stepsz

print stepsr,field[0][0],field[len(field)-1][0] 
hists=[]
fieldr=ROOT.TH2D("fieldr","Field R-Component (V/cm)",
                 stepsz,field[0][1],field[len(field)-1][1]+dz,
                 stepsr,field[0][0],field[len(field)-1][0]+dr)
setAxis(fieldr,"Z (cm)","Radius (cm)","Radial Field Strength (V/cm)")
if args.ntit:
    fieldr.SetTitle("")
hists.append(fieldr)

fieldz=ROOT.TH2D("fieldz","Field Z-Component (V/cm)",
                 stepsz,field[0][1],field[len(field)-1][1]+dz,
                 stepsr,field[0][0],field[len(field)-1][0]+dr)
setAxis(fieldz,"Z (cm)","Radius (cm)")
hists.append(fieldz)

pot=ROOT.TH2D("pot","Potential (V)",
              stepsz,field[0][1],field[len(field)-1][1]+dz,
              stepsr,field[0][0],field[len(field)-1][0]+dr)
setAxis(pot,"Z (cm)","Radius (cm)")
hists.append(pot)

absfield=ROOT.TH2D("absfield","Absolute Field (V/cm)",
              stepsz,field[0][1],field[len(field)-1][1]+dz,
              stepsr,field[0][0],field[len(field)-1][0]+dr)
setAxis(absfield,"Z (cm)","Radius (cm)")
hists.append(absfield)


if args.out!='':
    outname=args.out
else:
    outname=args.filename.replace('.txt','_out.txt')
    outname.replace('../','./')

print outname
outfile=open(outname,'w+')
outfile.write(str(nomEz)+' '
              +str(rmin)+' '+str(dr)+' '+str(stepsr)+' '
              +str(zmin)+' '+str(dz)+' '+str(stepsz)+'\n')


field.sort(key=itemgetter(0))


for i in range(len(field)):
    outfile.write(str(field[i][2])+' '+str(field[i][3])+'\n')
    fieldr.Fill(field[i][1]+dz/2,field[i][0]+dr/2,field[i][2])
    fieldz.Fill(field[i][1]+dz/2,field[i][0]+dr/2,field[i][3])
    pot.   Fill(field[i][1]+dz/2,field[i][0]+dr/2,field[i][4])
    absfield.Fill(field[i][1]+dz/2,field[i][0]+dr/2,math.sqrt(field[i][3]**2+field[i][2]**2))
        
print "Done filling histos"

outfile.close()

if args.batch and not args.pr:
    exit()

outfile=ROOT.TFile("field.root","recreate")
canvases=[]


title=args.filename.replace('.txt','')
title=args.prpath+title.split('/')[-1]
# {{{ Field R

c1=ROOT.TCanvas("c1","Field_R",20,20,900,400)
setMargin(c1)
canvases.append(c1)
fieldr.SetStats(0)
setRange(fieldr,-10,10)
if not args.batch:
    fieldr.Draw("colz")
c1.Update()
fieldr.Write()

# }}}
# {{{ Field Z
c2=ROOT.TCanvas("c2","Field_Z",930,0,900,400)
setMargin(c2)
canvases.append(c2)
fieldz.SetStats(0)
setRange(fieldz,args.zfieldRange[0],args.zfieldRange[1])
if not args.batch:
    fieldz.Draw("colz")
c2.Update()
fieldz.Write()
# }}}
# {{{ Pot
c3=ROOT.TCanvas("c3","Pot",20,500,900,400)
setMargin(c3)
canvases.append(c3)
pot.SetStats(0)
pot.SetContour(500)
c3.SetRightMargin(.2)
if not args.batch:
    pot.Draw("CONT1Z")
#pot.Draw("COLZ")
c3.Update()
pot.Write()
# }}}
# {{{ Field with arrows

c4=ROOT.TCanvas("c4","Field",930,500,900,400)
setMargin(c4)
canvases.append(c4)
absfield.SetStats(0)
setRange(absfield,args.zfieldRange[0],args.zfieldRange[1])
if not args.batch:
    absfield.Draw("colz")
arrows=[]
for i in range(len(field)):
    arrows.append(ROOT.TArrow())
    absl=math.sqrt(field[i][3]**2+field[i][2]**2)
    row=i%stepsz
    col=i-(row*stepsz)
    rowdiv=40
    coldiv=4
    flen=math.sqrt(rowdiv**2+coldiv**2)
    if absl>0:
        arrows[i].SetX1(field[i][1]+dr/2.)
        arrows[i].SetY1(field[i][0]+dr/2.)
        arrows[i].SetX2(field[i][1]+dr/2.-(field[i][3]/absl)*(dz/2)*flen)
        arrows[i].SetY2(field[i][0]+dr/2.-(field[i][2]/absl)*(dr/2)*flen)
        arrows[i].SetOption("|>")
        arrows[i].SetArrowSize(0.008)
        arrows[i].SetFillColor(1)
        if row%rowdiv==0 and col%coldiv==0:
            arrows[i].Draw()
 
c4.Update()
absfield.Write()

# }}}

if args.prel:
    for c in canvases:
        c.cd()
        drawPrelim()
        c.Update()
        

outfile.Close()
if args.pr:
    canvases[0].SaveAs(title+'.pdf[')
    for c in canvases:
        c.SaveAs(title+'.pdf')
    canvases[0].SaveAs(title+'.pdf]')
if not args.batch:
    u=raw_input("done?")