# HepSim validation script using Jas4pp or hs-toolkit 
# Task: Reconstruct pT amd Mjj cross sections using ProIO input file.
#       The particle records are saved as VarintPackedParticles 
#       The validation also uses weights from Pythia8  
# S.Chekanov (ANL), Aug. 2018 
#
from java.lang import *
from java.awt import Color,Font
from jhplot import  HPlot,P1D,HLabel,H1D
from jhplot.io import HBook
from hepsimproio import HepSim, FileMC 
from jhplot.utils import FileList
from hepsim import HepSim
from java.util import ArrayList
from hephysics.jet import ParticleD,JetN2
import sys

NMax=5 # max number of files for testing 
dataset="tev13pp_qcd_pythia8_proio"
tag="" # only truth level
url=""
if len(sys.argv)>1:
   flist=FileList.get(sys.argv[1],"proio")
   if (len(sys.argv)>2): NMax=int(sys.argv[2])
else:
   sites=HepSim.urlRedirector(dataset)
   url=sites[0]+"/"+tag
   flist=HepSim.getList(url)
if len(flist)==0: print "Error: No input file!"; sys.exit(0)
else: print "Reading "+str(len(flist))+" files. Nr of files= ",NMax

h1= H1D("PT(jet)",50,50,3000) 
h2= H1D("M(jj)",25,100,10000)
#h1.setFill(True)
#h1.doc()               # check its methods

ktjet=JetN2(0.4,"antikt",50) # E-mode, anti-KT,min pT=50 
# print ktjet.info()

nev=0;  Nfiles=len(flist)
if (NMax>-1 and len(flist)>Nfiles):       Nfiles=NMax           # limit to NMax files 
across,weightSumTot,lumi=0,0,0            # cross, lumi and total weights

um,lun=0,0 # varint encoding
mm=0
for m in range(Nfiles):             # loop over all files in this list    
   weightSum=0
   fmc=FileMC(url+flist[m])
   reader=fmc.reader()
   first=True
   while (reader.hasNext()):         # event loop 
      event=reader.next()
      metadata = event.getMetadata() # get metadata for this event 
      nev=nev+1
      if (nev%200==0):
           if (Nfiles==1): print "Event=",nev
           else: print "Event=",nev," done=",int((100.0*m)/Nfiles),"%"

      if (first):
        un=float(metadata["info:varint_energy"].toStringUtf8())
        lun=float(metadata["info:varint_length"].toStringUtf8())
        first=False

      if (reader.hasNext() == False):
          print "Metadata from last event:"
          #metadata=metadata.entrySet()
          print metadata["meta:description"].toStringUtf8();
          print metadata["meta:creation_time"].toStringUtf8()
          lumi=lumi+float(metadata["meta:luminosity_inv_pb"].toStringUtf8());
          across=across+float(metadata["meta:cross_section_pb"].toStringUtf8());
          print "-> Accumulated luminosity (pb-1) =",lumi
          print "-> Cross section (pb)=",float(metadata["meta:cross_section_pb"].toStringUtf8());
          # weight sum should be for entire file!
          weightSumTot=weightSumTot+weightSum;
          break

      # block with some MC parameters
      weight=1.0
      weightSum=0.0
      for entryID in event.getTaggedEntries("MCParameters"):  
           mc=event.getEntry(entryID)
           weight=mc.getWeight()
      for entryID in event.getTaggedEntries("Pythia8Parameters"): # Pythia8 specific  
           mc=event.getEntry(entryID)
           weightSum=mc.getWeightSum()

      # loop over particles 
      for entryID in event.getTaggedEntries("VarintPackedParticles"): 
           pa=event.getEntry(entryID)
           particles=ArrayList()
           for j in range(pa.getIdCount()): # loop over particles 
              if(pa.getStatus(j)==1):
                p=ParticleD(pa.getPx(j)/un,pa.getPy(j)/un,pa.getPz(j)/un,pa.getEnergy(j)/un)
                particles.add(p)
      ktjet.buildJets(particles)
      # ktjet.printJets()
      jets=ktjet.getJetsSorted() # return jets sorted in pT 
      if (len(jets)>0):
                 lead=jets[0]
                 h1.fill(lead.perp(),weight)
      if (len(jets)>1):
                 pt1=jets[0].perp()
                 pt2=jets[1].perp()
                 if (abs(jets[0].phi()-jets[1].phi())>2):
                         jets[0].add(jets[1])
                         mass=jets[0].mass()
                         h2.fill(mass,weight)

   mm=mm+1
   fmc.close()

cross=across/float(mm) # average cross section 
lumi=nev/cross;
print "Lumi=%.3e pb"%lumi
print "Total cross section (pb)=",cross
print "Total weight=",weightSumTot
h1.scale(cross/(h1.getBinSize()*weightSumTot))
h2.scale(cross/(h1.getBinSize()*weightSumTot))

c1 = HPlot("HepSim")
c1.setNameX("p_{T}(jet) [GeV]")
c1.setNameY("d σ / d p_{T} [pb / GeV]")
c1.visible(True)
c1.setMarginLeft(90)
c1.setAutoRange(True)
c1.cd(1,1)
c1.setRangeX(0,3000)
c1.setRangeY(0.00001,10000000)
c1.setLogScale(1,1)
c1.draw(h1)

l1=HLabel("Pythia8 QCD σ=%.3e pb"%cross, 0.5, 0.75, "NDC")
l1.setFont(Font("Helvetica", Font.PLAIN, 12))
c1.add(l1)

l2=HLabel("=HepSim=",0.8,0.85, "NDC")
l2.setColor(Color.gray)
l2.setFont(Font("Helvetica", Font.PLAIN, 14))
c1.add(l2)

# create file/image using name of the file
name="output"
if len(sys.argv[0])>0: name=sys.argv[0].replace(".py","")
file=HBook(name+".jdat","w"); print name+".jdat created"
file.write(h1)
file.close()
c1.export(name+".svg");    print name+".svg created"
# h1.toFile(name+".txt");  print name+".txt created"
# sys.exit(0)