# Example that shows how to make jets from PFA (reconstructed particles) using  LCIO format
# This example looks at jet resolution.
# This is a simple example that uses the Durham algorithm with pT>20 GeV
# It also does a Gaussian fit. 
# Use Z to bbar data 
#  Get 4 file in 2 thread (see HepSim manual):
#  > hs-get gev250ee_pythia6_qcd_all%rfull001 gev250ee_pythia6_qcd_all 2 4 
#
# author: S.Chekanov (ANL)


from org.lcsim.event import *
from org.lcsim.util import *
from hep.physics.vec import BasicHepLorentzVector,HepLorentzVector 
from java.util import *
from java.io import *
from org.lcsim.lcio import LCIOReader
from hep.io.sio import SIOReader
from hep.lcio.implementation.sio import SIOLCReader
from hep.lcio.implementation.io import LCFactory
from hep.lcio.event import *
from hep.lcio.io import *
from jhplot import *  # import graphics
from hephysics.particle import LParticle 
from hep.physics.jet import FixNumberOfJetsFinder
from hephysics.jet import ParticleD
from java.awt import *
import math 

# make list of files..
import glob
files=glob.glob("gev250ee_pythia6_qcd_all/*.slcio")
factory = LCFactory.getInstance()
reader = factory.createLCReader()
reader.open(files)

xmin=0.0
xmax=200.
h1=H1D("jet pt",100,xmin,xmax)
# h1.doc()  # look at API

# jet description
# http://jwork.org/dmelt/api/doc.php/hephysics/jet/JetN2
from hephysics.jet import JetN2 
ktjet=JetN2(0.5,"antikt",20)        # initialize jets with R=0.5, E-mode, anti-KT,min pT=20
print ktjet.info()                  # print its settings


nEvent=0
while(1):
     evt=reader.readNextEvent()
     if (evt == None): break
     nEvent=nEvent+1
     # print " file event: ",evt.getEventNumber(), " run=",evt.getRunNumber()
     if (nEvent%50==0): print "# Event: ",nEvent
     strVec = evt.getCollectionNames()
     if nEvent == 1:
            for col in  strVec:
                           print col           
     col = evt.getCollection("PandoraPFOCollection")
     nPFA = col.getNumberOfElements()
     # if (nEvent>2000): break

     alljets=[] # make a new list with jets 
     particles=ArrayList() # list of particles  
     for i in range(nPFA):
          # http://www.lcsim.org/sites/lcsim/apidocs/org/lcsim/lcio/SIOReconstructedParticle.html
          # http://www.lcsim.org/sites/lcsim/apidocs/org/lcsim/event/ReconstructedParticle.html
          pa=col.getElementAt(i)
          charge=pa.getCharge()
          p4=pa.getMomentum()  
          ee=pa.getEnergy() 
          typep=pa.getType() # type of PF 
          # http://jwork.org/dmelt/api/doc.php/hephysics/jet/ParticleD.html
          p=ParticleD(p4[0],p4[1],p4[2],ee);
          particles.add(p) # add particle to the list 
     ktjet.buildJets(particles)
     jets=ktjet.getJetsSorted()      # get a list with sorted jets
     if (len(jets)>0):
                 print "pT of a leading jet =",jets[0].perp()," GeV" 
                 h1.fill(jets[0].perp())

reader.close();

c1=HPlot("resolution")
# c1.doc() # look at API
c1.visible()
c1.setMarginLeft(80)
c1.setNameX("pT(jet, PFA) / pT(jet, truth)")
c1.setNameY("Events")
c1.setAutoRange()
c1.draw(h1) 
c1.export("mc_jets_antiKT.pdf")