import glob, sys, os, StringIO, ROOT

options = ["-pf","-f","-outfile","-draw", "-devfile","-pretty"]
opts = set(options)

#default parameters
pfilename = "tpc/parfiles/tpc.par"
pfilesname = "tpc/parfiles/tpc.files.par"
devfilename = "tpc/parfiles/DevMap_20MeV_plateau_mode1_mom3.672_thMin0.010.dat"


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]

infileGiven = 0
pretty = 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
    if arg == "-devfile":   #output file name
        devfilename = getArgsToNextOpt(sys.argv,iarg)[0]   
    if arg == "-f":   #input ROOT file name
        infilename = getArgsToNextOpt(sys.argv,iarg)[0]       
        infileGiven = 1
    if arg == "-pretty" :
        pretty = 1
        
        
if not infileGiven:
    print "No input ROOT file (containing the splines) given! Aborting..."
    exit()
        
#vdrift:
pfile = open(pfilename, "r")
pfilesfile = open(pfilesname, "r")

for line in pfile:
    if line.startswith("TpcGasFile") :
        gasline = (int)(getLast(line,' ').rstrip())
    if line.startswith("EField") :
        EField = (float)(getLast(line,' ').rstrip())
    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())      

lcount = 0
for line in pfilesfile:
    if lcount == gasline :
        gasfilename = line.rstrip()
    lcount+=1

print "Getting vDrift from",gasfilename
print "EField strength from parfile:",EField

ROOT.gROOT.ProcessLine(".x rootlogon.C") 
if pretty :
    ROOT.gROOT.ProcessLine(".x rootlogon_Bernhard_New.C")
    ROOT.gROOT.ProcessLine('gROOT->SetStyle("col")'); 
    ROOT.gROOT.ProcessLine('gROOT->ForceStyle()'); 
    
else:
     ROOT.gROOT.ProcessLine('gROOT->SetStyle("Plain")'); 
     ROOT.gROOT.ProcessLine('gStyle->SetPalette(1)'); 

gas = ROOT.TpcGas(gasfilename, EField)
print "Extracted vDrift:",gas.VDrift()

print
print "Initializing DevMap:"
devmap = ROOT.TpcDevmapCylWrapper(devfilename)


rFile = ROOT.TFile(infilename, "R")
tree = rFile.Get("cbmsim")

tree.SetBranchStatus("*", 0)
tree.SetBranchStatus("TpcLaserSplineFits.*", 1)
splines = {}
splineNames = { 0: 'spRecoRad', 1: 'spRecoPhi', 2: 'spCorrRad', 3 : 'spCorrPhi'}

for e in tree :
    count = 0
    for spl in e.TpcLaserSplineFits :
        name = splineNames[count]
        splines[name] = e.TpcLaserSplineFits.At(count)
        count += 1
        
splineRad = splines['spRecoRad']
splinePhi = splines['spRecoPhi']

nBinsZ = 250
nBinsR = 60

binWidthZ = (zMax - zMin) / nBinsZ
binWidthR = (rMax - rMin) / nBinsR

#histos
histSplineRad = ROOT.TH2D("SplineRecoRad", "SplineRecoRad", nBinsZ, zMin, zMax, nBinsR, rMin, rMax)
histDevmapRad = ROOT.TH2D("DevmapRad", "DevmapRad", nBinsZ, zMin, zMax, nBinsR, rMin, rMax)
histDiffmapRad = ROOT.TH2D("DevmapRadDiffMap", "DevmapRad Diff (reco - orig)", 
                           nBinsZ, zMin, zMax, nBinsR, rMin, rMax)
histDiffRad = ROOT.TH1D("DevmapRadDiff", "DevmapRad Diff (reco - orig)", 300,-0.11,0.11)
histDiffRad.GetXaxis().SetTitle('#xi_{rad, reco} - #xi_{rad} (cm)')

histSplinePhi = ROOT.TH2D("SplineRecoPhi", "SplineRecoPhi", nBinsZ, zMin, zMax, nBinsR, rMin, rMax)
histDevmapPhi = ROOT.TH2D("DevmapPhi", "DevmapPhi", nBinsZ, zMin, zMax, nBinsR, rMin, rMax)
histDiffmapPhi = ROOT.TH2D("DevmapPhiDiffMap", "DevmapPhi Diff (reco - orig)", 
                           nBinsZ, zMin, zMax, nBinsR, rMin, rMax)
histDiffPhi = ROOT.TH1D("DevmapPhiDiff", "DevmapPhi Diff (reco - orig)", 300,-0.11,0.11)
histDiffPhi.GetXaxis().SetTitle('#xi_{#phi, reco} - #xi_{#phi} (cm)')

for iz in range(nBinsZ) :
    z = zMin + iz*binWidthZ
    for ir in range(nBinsR) :
        r = rMin + ir*binWidthR
        vec = ROOT.TVector3(r,0,z)
        devRecoRad = splineRad.eval(z,r)
        devRecoPhi = splinePhi.eval(z,r)
        histSplineRad.SetBinContent(iz,ir,devRecoRad)
        result = devmap.value(vec)
        histDevmapRad.SetBinContent(iz,ir,result.X())
        histDiffmapRad.SetBinContent(iz,ir,devRecoRad-result.X())
        histDiffRad.Fill(devRecoRad-result.X())
        
        histSplinePhi.SetBinContent(iz,ir,devRecoPhi)
        histDevmapPhi.SetBinContent(iz,ir,result.Y())
        histDiffmapPhi.SetBinContent(iz,ir,devRecoPhi-result.Y())
        histDiffPhi.Fill(devRecoPhi-result.Y())
        
        

c1 = ROOT.TCanvas()
c1.Divide(2,2)
c1.cd(1)
histSplineRad.Draw("colz")
c1.cd(2)
histDevmapRad.Draw("colz")
c1.cd(3)
histDiffmapRad.Draw("colz")
c1.cd(4)
histDiffRad.Draw()

c1_pretty = ROOT.TCanvas("mh1","dsdsd",800,400)
histDiffRad.Draw()



c2 = ROOT.TCanvas()
c2.Divide(2,2)
c2.cd(1)
histSplinePhi.Draw("colz")
c2.cd(2)
histDevmapPhi.Draw("colz")
c2.cd(3)
histDiffmapPhi.Draw("colz")
c2.cd(4)
histDiffPhi.Draw()

c2_pretty = ROOT.TCanvas("meh2","dsdsdsd",800,400)
histDiffPhi.Draw()

input()