import ROOT, glob, math, sys, os
pandapath=os.environ.get('PANDAPATH')
sys.path.append(pandapath+'/macro/tpc/FOPI/python/argparse-1.2.1')
import argparse
from ROOT import std
from array import array

def readParfile(pFile):
    print "reading parameterfile", pFile," for alignment info"
    parFile=open(args.parf, "r")
    for lines in parFile:
        if lines.find("AlignmentFile:")!=-1:
            lines=lines.strip()
            line=int( (lines.split(":")[1]).split()[1] )
            TpcFiles=open("tpc/parfiles/tpc.files.par", "r")
            cTpcF=0
            for lin in TpcFiles:
                if cTpcF==line:
                    print "using Alignment file:",str(lin[:-1])
                    return str(lin[:-1])
                cTpcF+=1
    parFile.close()
    TpcFiles.close()

parser=argparse.ArgumentParser(description="starts the calculator for the deviation map.")
parser.add_argument("ef",help="file with the e-field",type=str)
parser.add_argument("parf",help='parameter file',type=str)
parser.add_argument("of",help='prefix for the outputfile',type=str)
parser.add_argument("--rBin",help='if not set all bins (nbinsr,nbinsz) will be calculated',default=-1,type=int)
parser.add_argument("--zBin",help='if not set all bins in Z will be calculated',default=-1,type=int)
parser.add_argument("--mom",help='momentum, dont know for what!',default=3.672,type=float)
parser.add_argument("--nr",help="number of bins in R to calculate",default=211,type=int)
parser.add_argument("--nz",help="number of bins in Z to calculate",default=1447,type=int)
parser.add_argument("--zStep",help='number of bins in one step',default=-1,type=int)
parser.add_argument("--specZ",help='calculate given z bins, komma separated list',default="-1",type=str)
parser.add_argument('--align',help='Aligment',action='store_true')
args = parser.parse_args()

if os.path.exists(args.ef):
    print "Using E-Field file:",args.ef
else:
    print "Bad E-Field file:",args.ef
    sys.exit(0)

if os.path.exists(args.parf):
    print "Using Parameter file:",args.parf
else:
    print "Bad Parameter file:",args.parf
    sys.exit(0)

if args.align:
    parFile=open(args.parf, "r")
    AlignFile=readParfile(args.parf)
    alignf=open(AlignFile, "r")
    count=0
    ar=[]
    found=0
    
    for line in alignf:
        if line.find("detName tpc")!=-1:
            found+=1
        if found>0:
            if line.find("translation")!=-1:
                trans=line.split()
                x=float(trans[1])
                y=float(trans[2])
                z=float(trans[3])
                found+=1
            if line.find("euler")!=-1:
                euler=line.split()
                phi  =float(euler[1])
                theta=float(euler[2])
                psi  =float(euler[3])
                found+=1
        if found==3:
            break
    print x, y, z, phi, theta, psi
    
    
ROOT.gROOT.ProcessLine(".x rootlogon.C") 
fRun = ROOT.FairRunSim()
fRun.SetBeamMom(args.mom)

efield = ROOT.TpcEFieldCyl(args.ef)
#efield.SetVerbose()
#bfield = ROOT.PndMultiField("FULL")
bfield = ROOT.FOPIField(0.616); 
bfield.SetTargetOffset(-3.5);
if args.align:
    bfield.SetAlignment(x,y,z,phi,theta,psi)

#bfield = ROOT.PndConstField()
#bfield.SetField(0.,0.,0.)
#bfield.SetFieldRegion(-50, 50,-50, 50, -2000, 2000);
#bfield.Init()

#the actual calculation
#efield=ROOT.TVector3(0,0,0)
#bfield=ROOT.TVector3(0,0,5)
drifter = ROOT.TpcInhFieldDrifter(efield, bfield,
                                  args.of, args.parf, args.nr, args.nz,
                                  args.rBin, args.zBin )
if args.zStep!=-1:
    drifter.setZStep(args.zStep)
if args.specZ!="-1":
    words=args.specZ.split(",")
    drifter.setManBins()
    for word in words:
        drifter.addZBin(int(word))
    
print "Starting the drifter...."
print "Writing Deviation Map output to",args.of

drifter.run()