import ROOT,argparse

def getFieldlineXZ(x0,z0,xmax,zmax,field):
    fieldline=ROOT.TPolyLine()
    x=x0
    z=z0
    pcount=0
    fieldline.SetPoint(pcount,z,x)
    while (abs(x)<xmax and abs(z)<zmax):
        pcount+=1
        oldx=x
        oldz=z
        if field.GetBx(oldx,0,oldz)==0 and field.GetBz(oldx,0,oldz)==0:
            break;
        x+=field.GetBx(oldx,0,oldz)/5.
        z+=field.GetBz(oldx,0,oldz)/5.
#        print "ox:",oldx,"oz:",oldz,"x:",x,"z:",z
        fieldline.SetPoint(pcount,z,x)
    return fieldline

ROOT.gROOT.ProcessLine(".x rootlogon.C") 
parser=argparse.ArgumentParser(description="creates a FOPI field and plots the field.")
parser.add_argument('--x',default=10)
parser.add_argument('--y',default=10)
parser.add_argument('--z',default=10)
parser.add_argument('--xoff',default=0)
parser.add_argument('--yoff',default=0)
parser.add_argument('--zoff',default=0)
parser.add_argument('--tar',default=0)

args = parser.parse_args()

bfield=ROOT.FOPIField(0.616)
bfield.SetTargetOffset(float(args.tar))
x=int(args.x)
y=int(args.y)
z=int(args.z)
zoff=float(args.zoff)

hxyBz=ROOT.TH2D('hxyBz',"Bz Field in xy plane",2*x,-x,x,2*y,-y,y)

hxzBz=ROOT.TH2D('hxzBz',"Bz Field in xz plane",2*z,-z,z,2*x,-x,x)
hxzBx=ROOT.TH2D('hxzBx',"Bx Field in xz plane",2*z,-z,z,2*x,-x,x)

hxyBz.SetStats(0)

hxzBz.SetStats(0)
hxzBx.SetStats(0)

for i in range(-x,x):
    for j in range(-y,y):
        hxyBz.Fill(i,j,bfield.GetBz(i,j,zoff))
    for j in range(-z,z):
        hxzBz.Fill(j,i,bfield.GetBz(i,0,j+zoff))
        hxzBx.Fill(j,i,bfield.GetBx(i,0,j+zoff))

c1=ROOT.TCanvas()
hxyBz.Draw("colz")
c2=ROOT.TCanvas()
hxzBz.Draw("colz")
line=[]
for i in range(-x,x,1):
    line.append(getFieldlineXZ(i,0,x,z,bfield))
    line[-1].Draw()
c2.Update()
print "fieldlines done"
c3=ROOT.TCanvas()
hxzBx.Draw("colz")

u=raw_input()