import argparse, ROOT, math

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('filename2',help='name of the COMSOL field file for add in')
parser.add_argument('--cutlz',help='lower value to cut and past z',default=7)
parser.add_argument('--cutuz',help='lower value to cut and past z',default=13)
parser.add_argument('--filllz',help='lower value to cut and past z',default=16)
parser.add_argument('--filluz',help='lower value to cut and past z',default=17.6)
args = parser.parse_args()

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

r=[]
z=[]
Er=[]
Ez=[]

radd=[]
zadd=[]

rmin=1000
zmin=1000

for line in file:
    words=line.split()
    if words[0]=='%':
        continue
    r.append(float(words[0]))
    z.append(float(words[1]))
    Er.append(float(words[2])/100.)
    Ez.append(float(words[3])/100.)
    if float(words[0])<rmin:
        rmin=float(words[0])
    if float(words[1])<zmin:
        zmin=float(words[1])
for line in file2:
    words=line.split()
    if words[0]=='%':
        continue
# dreckiges miststueck, von 16-17.6 schneiden
    if float(words[1])>args.cutlz and float(words[1])<args.cutuz:
#        print float(words[2])/100.,float(words[0]),float(words[1])
        if float(words[0])<51 or float(words[0])>149:
            radd.append(float(words[2])/100.)
        else:
            radd.append(0.)

        zadd.append(float(words[3])/100.)
 
print "done reading the files"       
#print "len(r)",len(r),"lenz",len(z),r[0],r[len(r)-1
field=ROOT.TH2D("field","field",int(len(r)/1000),r[0],r[len(r)-1],int(len(z)/1000),z[0],z[len(z)-1])
fieldz=ROOT.TH2D("fieldz","fieldz",int(len(r)/1000),r[0],r[len(r)-1],int(len(z)/1000),z[0],z[len(z)-1])


dr=abs(r[0]-r[1])
dz=dr#abs(z[0]-z[1])
nomEz=309.6
stepsr=abs(r[0]-r[len(r)-1])/dr
stepsz=abs(z[0]-z[len(z)-1])/dz

newr=[]
newz=[]
newEr=[]
newEz=[]

counter=-1
zoff=0
rcounter=0

for i in range(len(r)):
# dont forget to loop over z too, now they are disentangled!!!!
    counter+=1
    newr.append(r[i])
    newz.append(z[i]+zoff)
    newEr.append(Er[i])
    newEz.append(Ez[i])

    if z[i]>=args.filllz and r[i]==50. and z[i]<args.filluz:
#        print "now",r[i],z[i]
        while newz[counter]<71.9:
            for j in range(len(radd)):
#                dont forget to loop over the z also
                counter+=1
#                print newr[counter-1],counter,rcounter,i
                newr.append(r[i+rcounter])

                if newr[counter]==50.:
#                    print 'bla'
                    zoff+=.1
                    rcounter=0

                newz.append(z[i]+zoff)
#                print newz[counter]
                newEz.append(zadd[j])
#                print zadd[j]
                newEr.append(radd[j])
              
                rcounter+=1
#                print counter, newz[counter], newr[counter], rcounter,r[i+rcounter]
#                u=raw_input()

print "done creating fields"
outfile=open("tpcfield.txt",'w+')
outfile.write(str(nomEz)+' '
              +str(rmin)+' '+str(dr)+' '+str(stepsr)+' '
              +str(zmin)+' '+str(dz)+' '+str(stepsz)+'\n')

print len(z),len(newz),newz[0],newz[len(newz)-1]
newfield=ROOT.TH2D("newfield","newfield",int(len(newr)/1000),newr[0],newr[len(newr)-1],int(len(newz)/1000),newz[0],newz[len(newz)-1])
newfieldz=ROOT.TH2D("newfieldz","newfieldz",int(len(newr)/1000),newr[0],newr[len(newr)-1],int(len(newz)/1000),newz[0],newz[len(newz)-1])

for i in range(len(newz)):
    if newfield.GetBinContent(newfield.FindBin(newr[i],newz[i]))==0:
        newfield.Fill(newr[i],newz[i],newEr[i])
    if newfieldz.GetBinContent(newfieldz.FindBin(newr[i],newz[i]))==0:
        newfieldz.Fill(newr[i],newz[i],newEz[i])


for i in range(len(r)):
    outfile.write(str(Er[i])+' '+str(Ez[i])+'\n')
    if field.GetBinContent(field.FindBin(r[i],z[i]))==0:
        field.Fill(r[i],z[i],Er[i])
    if fieldz.GetBinContent(fieldz.FindBin(r[i],z[i]))==0:
        fieldz.Fill(r[i],z[i],Ez[i])

print "Done filling histos"

field.SetStats(0)
field.GetZaxis().SetRangeUser(-100,100)
field.Draw("colz")
c2=ROOT.TCanvas()
newfield.SetStats(0)
newfield.GetZaxis().SetRangeUser(-100,100)
newfield.Draw("colz")
c3=ROOT.TCanvas()
fieldz.SetStats(0)
fieldz.GetZaxis().SetRangeUser(0,500)
fieldz.Draw("colz")
c4=ROOT.TCanvas()
newfieldz.SetStats(0)
newfieldz.GetZaxis().SetRangeUser(0,500)
newfieldz.Draw("colz")
u=raw_input()