# Extract FOPI MC start values from fopiroot SHIELD simulation
# Use --plot option to plot graphs
# Author: Felix Boehmer

import argparse, ROOT, os, glob, array, math, datetime

parser=argparse.ArgumentParser(description='Extract FOPI MC start values from fopiroot SHIELD simulation')
parser.add_argument('files', metavar='files', type=str, nargs='+',
                    help='a input file (list)')
parser.add_argument('--plot',help="Calculate some QA and plot",action='store_true')

args = parser.parse_args()

if args.files == None or len(args.files) < 1 :
    print "No input files given. Aborting"
    exit()



ROOT.gROOT.ProcessLine(".x $VMCWORKDIR/gconfig/rootlogon.C") 


rfiles = []
matchHist = ROOT.TH1D("matchHist","Overlap TPC-Shield",5,0,5)
dxHist = ROOT.TH1D("dxHist","Spatial step TPC - Shield",300,0,5)
dEHist = ROOT.TH1D("dEHist","Energy step TPC - Shield",300,0,100)
dEHist.GetXaxis().SetTitle("Energy Loss (MeV)")


for f in sorted(args.files) :
    print f
    rfile = ROOT.TFile(f,"read")
    rfiles.append(rfile)
    tree = rfile.Get("cbmsim")
    
    if tree == None :
        print "   ... Skipping this one ... "
        continue
    
    tree.SetBranchStatus("*", 0)
    tree.SetBranchStatus("MCTrack.*", 1)
    tree.SetBranchStatus("TpcPoint.*", 1)
    tree.SetBranchStatus("TpcShieldPoint.*", 1)
    
    #enable caching
    tree.AddBranchToCache("*", True)    
    tree.SetCacheLearnEntries(1)

    outname = f.replace(".root",".dat")
    outf = open(outname,"write")
    outf.write("Converted from %s \nOn %s\n"%(f,datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S")))
    outf.write("Units: Momentum (GeV/c), Position (cm)\n\n")
    ec = 0

    for e in tree :
        matchedDict = {} 
        writeDict = {}
        trcount = 0
        for tr in e.MCTrack :
            #check for TPC and SHIELD (=kSTT) points:
            ntpc = tr.GetNPoints(10)
            nshld = tr.GetNPoints(9)
            matched = 0
            write = 0
            if ntpc > 0 and nshld > 0 :
                matchHist.Fill(1)
                matched = 1
                write = 1
            elif ntpc > 0:
                matchHist.Fill(0)
            elif nshld > 0 :
                write = 1
            matchedDict[trcount] = matched
            writeDict[trcount] = write
            trcount += 1

            
        tpDict = {}    
        for tp in e.TpcPoint :
            trID = tp.GetTrackID()
            if trID not in tpDict.keys() :
                tpDict[trID] = []
            tpDict[trID].append(tp)    
            
        tspDict = {}        
        for tsp in e.TpcShieldPoint :
            trID = tsp.GetTrackID()
            if trID not in tspDict.keys() :
                tspDict[trID] = []
            tspDict[trID].append(tsp)    

        #write
        outf.write("--------- EVENT %d ------------------------\n"%ec)
        for key, val in writeDict.iteritems() :
            if val == 1 :
                out = "     TrackID: %4.d  PDG: %10.d   4-Mom: (%.8f,%.8f,%.8f,%.8f)   Pos: (%.8f,%.8f,%.8f)   \n"
                PDG = e.MCTrack.At(key).GetPdgCode()
                mass = e.MCTrack.At(key).Get4Momentum().M()
                firstPx = tspDict[key][0].GetPx()
                firstPy = tspDict[key][0].GetPy()
                firstPz = tspDict[key][0].GetPz()
                firstX = tspDict[key][0].GetX()
                firstY = tspDict[key][0].GetY()
                firstZ = tspDict[key][0].GetZ()
                first = ROOT.TVector3(firstX,firstY,firstZ)
                firstP = ROOT.TVector3(firstPx,firstPy,firstPz)
                energy = math.sqrt(mass**2 + firstP.Mag()**2)
                outf.write(out%(key,PDG,energy,firstPx,firstPy,firstPz,firstX,firstY,firstZ))
                    

       
            
                
        if args.plot :         
            for key, val in matchedDict.iteritems() :
                if val == 1 :
                    PDG = e.MCTrack.At(key).GetPdgCode()
                    #if not e.MCTrack.At(key).IsGeneratorCreated() :
                    #    continue
                    #if PDG != 211 :
                    #    continue
                
                
                    lastX = tpDict[key][-1].GetX()
                    lastY = tpDict[key][-1].GetY()
                    lastZ = tpDict[key][-1].GetZ()
                    
                    firstX = tspDict[key][0].GetX()
                    firstY = tspDict[key][0].GetY()
                    firstZ = tspDict[key][0].GetZ()
                
                    last = ROOT.TVector3(lastX,lastY,lastZ)
                    first = ROOT.TVector3(firstX,firstY,firstZ)
                    dx = last-first
                    dxHist.Fill(dx.Mag())
                
                    #get mass:
                    mass = e.MCTrack.At(key).Get4Momentum().M()

                    #calculate energy loss
                    lastPx = tpDict[key][-1].GetPx()
                    lastPy = tpDict[key][-1].GetPy()
                    lastPz = tpDict[key][-1].GetPz()
                    lastP = ROOT.TVector3(lastPx,lastPy,lastPz)

                    firstPx = tspDict[key][0].GetPx()
                    firstPy = tspDict[key][0].GetPy()
                    firstPz = tspDict[key][0].GetPz()
                    firstP = ROOT.TVector3(firstPx,firstPy,firstPz)
                    
                    eLast =  math.sqrt(mass**2 + lastP.Mag()**2)
                    eFirst =  math.sqrt(mass**2 + firstP.Mag()**2)
                    
                    dEHist.Fill((eLast-eFirst)*1.e3)
                
        ec += 1


if args.plot :
    canv1 = ROOT.TCanvas("canv1","",800,400)
    matchHist.Draw()

    canv2 = ROOT.TCanvas("canv2","",800,800)
    canv2.Divide(1,2)
    canv2.cd(1)
    dxHist.Draw()
    canv2.cd(2)
    dEHist.Draw()
    raw_input()   

print "Outfile written to",outname