from java.awt import Color
from java.util import ArrayList
from jhplot  import *
from hephysics.jet import *
from hephysics.particle import *
from jhplot.shapes import *

plist=ArrayList()
# add px,py,pz,e
p=ParticleD(-0.880,-1.233,-157.431,157.439); plist.add(p)
p=ParticleD(10.0, -2.3211,-15.0785, 18.250); plist.add(p)
p=ParticleD(552.60,-143.38,127.405, 584.95); plist.add(p)
p=ParticleD(-58.56, 13.672,-60.729, 85.479); plist.add(p)
p=ParticleD(249.62, 59.98,-252.05,  359.78); plist.add(p)
p=ParticleD(8.31, -1.26, -13.692,   16.072); plist.add(p)
p=ParticleD(-132.34, 31.19,-133.81, 190.7);  plist.add(p) 
p=ParticleD(31.629, -7.989, 7.4822, 33.47);  plist.add(p) 

R=0.6 # jet radius. Use anti-kT jets 
ktjet = SCJet(R, 1, -1, 200.0,True);
ktjet.buildJets(plist);
ktjet.printJets()
jets=ktjet.getJetsSorted()

c1 = HPlot("Canvas",600,600)
c1.setGTitle("anti-kT jets")
c1.visible(1)
c1.setRange(-4,4,-7,7)
c1.setNameX("η")
c1.setNameY("φ [rad]")

pp=P1D("particles")

for i in range(plist.size()):
      p=plist.get(i)
      eta=p.getRapidity()
      phi=p.getPhi()
      pp.add(eta,phi)
      # psize=int(p.getEt()/10.)
      # pp.setSymbolSize(2)

# show jets
for k in range(jets.size()):
    jet=P1D("jet Nr "+str(k))
    jet.setColor(Color.red)
    jet.setSymbolSize(7)
    jj=jets.get(k)
    etaJ=jj.getRapidity()   
    phiJ=jj.phi()
    jet.add(etaJ,phiJ)
    c1.draw(jet)
    cic= Circle(etaJ, phiJ, R)
    cic.setFill(0)
    cic.setColor(Color.red)
    cic.setTransparency(0.5)
    c1.add(cic)

c1.draw(pp)