# Example that shows how to make jets from PFA (reconstructed particles) using  LCIO format
# This example makes invariant mass of 2 reconstructed jets and makes Z peak
# 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_zpole_bbar%rfull001 gev250ee_pythia6_zpole_bbar  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 java.awt import *
from math import *

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

h1=H1D("M(jj)",50,40,140)
# h1.doc()  # look at API

# jet description
# http://java.freehep.org/freehep-physics/apidocs/hep/physics/jet/package-summary.html
# this is Jade
from hep.physics.jet import DurhamJetFinder,FixNumberOfJetsFinder
# fjet=FixNumberOfJetsFinder(2) # request 2 jets. The Jade algorithm at work 
fjet=DurhamJetFinder(0.05)

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()

     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 
          p=BasicHepLorentzVector(ee,p4[0],p4[1],p4[2])
          particles.add(p) # add particle to the list 
     #print particles.size()
     if (particles.size()>1): # ask for >1 particles
       fjet.setEvent(particles)
       # print "Nr of jets=", fjet.njets()  
       alljets=[] # make a new list with jets 
       for i in range(fjet.njets()):
         # print "Jet=",i," Nr of particles in jet =",fjet.nParticlesPerJet(i) 
         pjet=fjet.jet(i)   # make  HepLorentzVector
         ee=pjet.t()  # magnitude
         p3=pjet.v3() # 3-vector  
         #print " ",ee,p3
         pl=LParticle(p3.x(),p3.y(),p3.z(),ee) # convert to a class with convinient kinematic methods  
         if (pl.perp()>20): alljets.append(pl)  
       if (len(alljets)>1): # calculate invariant mass of two jets 
         for i1 in range(0,len(alljets)-1):
          p1=alljets[i1]
          phi1=p1.phi()
          for i2 in range(i1+1,len(alljets)):
              p2=alljets[i2]
              phi2=p2.phi()
              p1.add(p2)
              if (abs(phi1-phi2)>1.5):
                mass=p1.calcMass()
                h1.fill(mass)
reader.close();

xmin=70 
xmax=110 
f=HFitter("chi2")   # build fitter 
f.setFunc("g")      # gaussian + constant  
func=f.getFunc()
#print func.parameterNames()
f.setPar('mean',90); f.setPar('amplitude',100)
f.setPar('sigma',8)
f.setRange(xmin,xmax) # fit range 
f.fit(h1)
fresult=f.getFittedFunc() # return final function 
res=f.getResult()
fPars  = res.fittedParameters()
fPars     = res.fittedParameters()
fParErrs  = res.errors()
fParNames = res.fittedParameterNames()
print "Fit results:"
for i in range(fresult.numberOfParameters()):
   print "par=",i,"=", fParNames[i]+" : "+str(fPars[i])+" +- "+str(fParErrs[i])

p1=P1D(h1)   # convert histogram to P1D array
# p1.doc()   # show its methods


c1=HPlot()
# c1.doc() # look at API
c1.visible()
c1.setMarginLeft(90)
c1.setNameX("M(jj) [GeV]")
c1.setNameY("Events")
c1.setRange(40,140,0,150)
# c1.draw(h1)
# c1.drawStatBox(h1)
c1.draw(p1)

# make F1D function from IFunction in the same range as histogram
f2 = F1D("Gauss fit",fresult,xmin,xmax)
f2.setColor(Color.blue)
f2.setPenWidth(2);
c1.draw(f2) # draw fit line  

q1="mean=%3.1f ± %3.1f GeV"%(fPars[1],fParErrs[1]) 
q2="σ=%3.1f ± %3.1f GeV"%(fPars[2],fParErrs[2]) 

l2=HLabel(q1,0.7,0.7, "NDC")
l2.setColor(Color.blue)
l2.setFont(Font("Helvetica", Font.PLAIN, 16))
c1.add(l2)

l2=HLabel(q2,0.7,0.64, "NDC")
l2.setColor(Color.blue)
l2.setFont(Font("Helvetica", Font.PLAIN, 16))
c1.add(l2)

c1.export("mc_jets.pdf")