# Validation script for HepSim using ProIO file format.
# Task: It plots the energy of scattered elecrons and the integrated luminosity 
# Run this script using Jas4pp or hs-toolkit 
# @author: S.Chekanov (ANL)

from java.lang import *
from hepsimproio import HepSim, FileMC 
from jhplot.utils import FileList
from java.awt import Color,Font
from jhplot import  HPlot,HPlot3D,P1D,HLabel,H2D,H1D
from jhplot.io import HBook
from hephysics.particle import LParticle 
import sys

Nmax=5 # max number of files 
dataset="gev35ep_pythia8_dis1q2ct14lo"
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)

h1=H1D("E(ele)",50,0,10)
lumi=0; nev=0;  # Nfiles=len(flist)
Nfiles=len(flist)
if (Nmax>0): Nfiles=Nmax            # limit the number of files
for m in range(Nfiles):             # loop over all files in this list    
   print flist[m]
   fmc=FileMC(url+flist[m])
   reader=fmc.reader()
   while (reader.hasNext()):             # event loop 
      event=reader.next()
      metadata = event.getMetadata() # get metadata for this event 
      # last event has useful information on cross sections, log file etc..
      if (reader.hasNext() == False):
          print "Metadata from last event:"
          #metadata=metadata.entrySet()
          print metadata["meta:description"].toStringUtf8();
          print "Cross section (pb)=",float(metadata["meta:cross_section_pb"].toStringUtf8());
          print metadata["meta:creation_time"].toStringUtf8()
          lumi=lumi+float(metadata["meta:luminosity_inv_pb"].toStringUtf8());
          print "-> Accumulated luminosity (pb-1) =",lumi
          #for entry in  metadata.entrySet(): print all if you need 
          #      skey=entry.getKey();
          #      name=entry.getValue().toStringUtf8();
          #      print skey,name    # print metadata if you need 
          break

      nev=nev+1
      if (nev%1000==0):
            print "Event=",nev," done=",int((100.0*m)/Nfiles),"%"
      for entryID in event.getTaggedEntries("Particle"): # particle loop  
            par=event.getEntry(entryID)
            v=par.getVertex()
            p=par.getP()
            pdg = par.getPdg()
            charge=int(par.getCharge()/3.) 
            status=par.getStatus()
            mass =par.getMass()
            px,py,pz = p.getX(),p.getY(),p.getZ()         # px,py,pz 
            x,y,z,t = v.getX(),v.getY(),v.getZ(),v.getT() # vertex components  
            ee = par.getEnergy()
            # print par # print this particle
            p=LParticle(px,py,pz,ee,mass) # also use this to calculate eta,rapidity,phi etc.  
            p.setCharge(charge)
            if (abs(pdg)==11 and status==1): # get scattered electron and exit event  
                                     h1.fill(ee);
                                     break  
   reader.close()

print "# Total events processed=",nev
print "# Total Luminosity=",lumi," pb-1" 

c1=HPlot() # plotting part
c1.visible()
c1.setAutoRange()
c1.setMarginLeft(90)
c1.setNameX("E_{ele} [GeV]")
c1.setNameY("Events")
c1.setRangeX(0,10)
c1.draw(h1)

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

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"
# sys.exit(0)