# HepSim script for hs-tools or Jas4pp 
# Part of =HepMC= : https://atlaswww.hep.anl.gov/asc/hepsim/
# Calculates particle masses from the truth level for the leading jet
# S.Chekanov (ANL)
from java.awt import Color,Font
from java.lang import *
from proto import FileMC            
from jhplot import HPlot,H1D,P1D,HLabel  # graphics
from jhplot.utils import FileList
from hephysics.particle import LParticle 
from hephysics.hepsim import PromcUtil
from hephysics.jet import JetN2 
from hepsim import HepSim
import sys,math 

ktjet=JetN2(0.5,"antikt",500) # anti-KT,E-mode, R=0.5,min pT=50 GeV,fast
# print ktjet.info()

url="";  TotalEvents=20000
default_www="http://mc.hep.anl.gov/asc/hepsim/events/mumu/5tev/mm5tev_pythia6_zprime5tev_ttbar/"
# default_www="http://mc.hep.anl.gov/asc/hepsim/events/mumu/40tev/mm40tev_pythia6_zprime40tev_qq/"
if len(sys.argv)>1: 
   if sys.argv[1].startswith("http"): 
                            flist=HepSim.getList(sys.argv[1])
                            url=sys.argv[1]
   else: flist=FileList.get(sys.argv[1],"promc")
   if (len(sys.argv)>2): TotalEvents=int(sys.argv[2]) 
else:
   url=default_www; flist=HepSim.getList(url)
if len(flist)==0: print "Error: No input file!"; sys.exit(0)
else: print "Reading "+str(len(flist))+" files. Nr events= ",TotalEvents


h1= H1D("PT(lead jet)",50,0,50000)  
h2= H1D("Particle mass",500,0,2)
h1.setFill(True)
h2.setFill(True)
#h1.doc()               # check its methods

c1 = HPlot("=HepSim=",500,500)
c1.visible(True)
c1.setAutoRange(True)
c1.setMarginLeft(90)
c1.setLegend(0)

cross=0; nev=0;  Nfiles=len(flist)
for m in range(Nfiles):              # loop over all files in this list    
   file=FileMC(url+flist[m])         # get input file
   header = file.getHeader()
   un=float(header.getMomentumUnit()) # conversion units
   lunit=float(header.getLengthUnit())
   if m==0:
       print "ProMC v=",file.getVersion(), "M unit=",un,"L unit=",lunit 
   if (nev>TotalEvents): print "Max Nr of events are done"; break # stop after some events
   for i in range(file.size()):
      if (nev>TotalEvents): break
      nev=nev+1
      if (nev%500==0):
           if (Nfiles==1): print "Event=",nev
           else: 
                print "Event=",nev," done=",int((100.0*m)/Nfiles),"%"
      if (nev%2000==0):
                 c1.clearData()
                 #c1.setNameX("p_{T}(lead. jet) [GeV]")
                 #c1.setNameY("Entries")
                 #c1.draw(h1)
                 c1.draw(h2)
                 print "Update canvas"
      eve = file.read(i)
      pa = eve.getParticles()    # particle information
      allpart=[]
      for j in range(pa.getPxCount()):
         if pa.getStatus(j)==1:
            if (abs(pa.getPdgId(j))==12 or abs(pa.getPdgId(j))==14 or  abs(pa.getPdgId(j))==15): continue # neutrino 
            par=LParticle(pa.getPx(j)/un,pa.getPy(j)/un,pa.getPz(j)/un,pa.getEnergy(j)/un,pa.getMass(j)/un)
            if (par.perp()>0.5):
                  allpart.append(par)

      #h2.fill( pa.getMass() )
      # ve = eve.getEvent()       # event information
      particles=PromcUtil.getParticleDList(file.getHeader(), pa, 1, -1, 1000);
      ktjet.buildJets(particles);
      # ktjet.printJets();
      jets=ktjet.getJetsSorted()
      if len(jets)>1:
        for jj in range(2):
               eta=jets[jj].getRapidity()
               phi=jets[jj].getPhi()
               h1.fill(jets[jj].perp())
               for pp in allpart:
                      phiL=pp.phi()
                      etaL=pp.rapidity()
                      dR2=(phi-phiL)*(phi-phiL) + (eta-etaL)*(eta-etaL)
                      dR=math.sqrt(dR2) 
                      if (dR<0.5):
                               h2.fill( pp.getMass() ) 

                 #h1.fill(jets[0].perp())
                 #for index in xlist:
                 #    part=particles[index]
                     #print part.mass() 
                     #print  part.mass(), type(part) 
                     #h2.fill( part.mass() )

   stat = file.getStat()
   cross=stat.getCrossSectionAccumulated()
   erro=stat.getCrossSectionErrorAccumulated();
   file.close()

lumi=nev/cross;
print "Lumi=%.3e pb"%lumi
print "Total cross section (pb)=",cross
hNew=h1.getDividedByBinWidth()
hNew.scale(1.0/lumi)
xsec=P1D(hNew)
xsec.setErr(1)       # show errors
xsec.setColor(Color.blue)

# show final cross section
c1.clearData()
c1.clear()
# c1.setLogScale(1,1)
#c1.setRangeY(0,01)
c1.setRangeX(0,2)
c1.setMarginLeft(90)
#c1.setNameX("p_{T}(lead. jet) [GeV]")
#c1.setNameY("d &sigma; / d p_{T} [pb/GeV]")
c1.setNameX("particle mass [GeV]")
c1.setNameY("Events")

# c1.draw(xsec)
# h1.toTable() 
c1.draw(h2)


# l1=HLabel("&sigma;=%.3e pb"%cross, 0.4, 0.75, "NDC")
l1=HLabel("<m>=%.3e GeV"%h2.mean(), 0.4, 0.75, "NDC")
l1.setFont(Font("Helvetica", Font.PLAIN, 15))
c1.add(l1)

l2=HLabel("=HepSim=",0.6,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","")
from jhplot.io import HBook 
file=HBook(name+".jdat","w"); print name+".jdat created"
file.write(xsec)
file.write(h2)
file.close()
c1.export(name+".svg");    print name+".svg created"
# xsec.toFile(name+".txt");  print name+".txt created"
# sys.exit(0)