import glob, sys, os, ROOT

options = ["-pf","-idens","-outfile","-draw"]
opts = set(options)

#default parameters
pfilename = "tpc/parfiles/tpc.par"
outfilename = "LaserGridVertical.dat"
iondens = 45
width = 0.03
draw = 0       #interactive mode

def getArgsToNextOpt(stringlist, index) :
    id = index+1
    args = []
    while(id < len(stringlist)):
        if stringlist[id] in options:
            return args
        args.append(stringlist[id])
        id+=1
    return args    

def getLast(pline,split) :
    chunks = pline.split(split)
    return chunks[len(chunks)-1:][0]


for iarg in range(len(sys.argv)) :
    arg = sys.argv[iarg]
    if arg == "-pf":   #tpc parameter file
        pfilename = getArgsToNextOpt(sys.argv,iarg)
    if arg == "-idens":   #laser track iondensity (cm^-1)
        iondens = (float) (getArgsToNextOpt(sys.argv,iarg))
    if arg == "-outfile":   #output file name
        outfilename = getArgsToNextOpt(sys.argv,iarg)[0]        
    if arg == "-draw" :
        draw  = 1
        

DIR = 'd'

#geometry:
pfile = open(pfilename, "r")

for line in pfile:
    if line.startswith("zGem") :
        zMin = (float)(getLast(line,' ').rstrip())
    if line.startswith("zMax") :
        zMax = (float)(getLast(line,' ').rstrip())
    if line.startswith("rMin") :
        rMin = (float)(getLast(line,' ').rstrip())
    if line.startswith("rMax") :
        rMax = (float)(getLast(line,' ').rstrip())    
        
pfile.close()
print "writing to",outfilename


outfile = open(outfilename, "w")
#tracks along drift : nTracks tracks equally dist. along radial direction, OFFSET away from inner/outer fieldcage
# one row every ANGDIST degrees, starting at startAng degrees. Repeat nRows times

nTracks = 6
nRows = 8
startAng = 20
ROOT_ANGLE_OFFSET =   3./2. * ROOT.TMath.Pi() #start at -1*yaxis
ANGDIST = 20
OFFSET = 0.5


rDist = (rMax - rMin - 2*OFFSET)/(nTracks-1)
angles = [startAng+i*ANGDIST for i in range(nRows)]
radii = [OFFSET+i*rDist for i in range(nTracks)]
dirV = ROOT.TVector3(0.,0.,-1.)

print radii
print angles

preamble = "#This file was created using macro createLaserFileNew.py\n\n" 
preamble += "#FORMAT: V id startX startY startZ VecX VecY VecZ iondens beamwidth time\n" 
preamble += "#where 'V' can either be 'd' or 'e' declaring Vec as direction or end-position, respectively\n" 
preamble += "#iondens : (nIons/cm)\n" 
preamble += "#width   : (cm)\n" 
preamble += "#time    : (s)\n" 
preamble += "#and id has to run continuously from 0 to idMax due to TClonesArray restrictions\n\n\n" 

outfile.write(preamble)

time = 0
count = 0 
polyMap = {}
step = 0.5
stepV = ROOT.TVector3(step*dirV.X(),step*dirV.Y(),step*dirV.Z())
for ang in angles : 
    for rad in radii :
        start = ROOT.TVector3(rad+rMin,0,zMax-0.2)
        start.SetPhi(ang*ROOT.TMath.Pi()/180. + ROOT_ANGLE_OFFSET)
        out = '%c %i %f %f %f %f %f %f %f %f %i \n' % (DIR,count,start.X(),start.Y(),start.Z(),
                                                       dirV.X(),dirV.Y(),dirV.Z(),
                                                       iondens,width,time)
        outfile.write(out)
        count+=1
        #plotting part
        if draw:
            poly = ROOT.TPolyLine3D(2)
            poly.SetPoint(0,start.X(),start.Y(),start.Z())
        #determine endpoint:
            end = ROOT.TVector3(start)
            while(end.Perp() > rMin and end.Perp() < rMax and 
                  end.Z() > zMin and end.Z() < zMax) :
                end += stepV
            poly.SetPoint(1,end.X(),end.Y(),end.Z())              
            polyMap[count] = poly

if draw:
#plot
    canv = ROOT.TCanvas()
    for (laserid, poly) in polyMap.items() :
        poly.Draw()

    input()

outfile.close()