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

plist=ArrayList()
def buildCluster(N,R,c_eta,c_phi,energy):
   r = Random()
   for i in range(N):
       pt=energy*abs(r.nextGaussian())
       phi=c_phi+2*R*(r.nextDouble()-0.5)
       eta=c_eta+2*R*(r.nextDouble()-0.5)
       p=ParticleD(); p.setPtEtaPhiM(pt,eta,phi,0)
       plist.add(p)

R=1.0
buildCluster(200,3.14,0,0,1)  # many random particles
buildCluster(100,R/2,1,1,10)  # 10 particles at eta=1 and phi=1
buildCluster(80,R/2,-1,-2,20) # 5 particles at eta=-1 and phi=-2

ktjet =JetN2(R,"antikt",20)   # min pT=20 GeV 
ktjet.buildJets(plist)
ktjet.printJets()
jets=ktjet.getJetsSorted()

c1 = HPlot("Canvas",500,500)
c1.visible()
c1.setRange(-4,4,-5,5)
c1.setNameX("η")
c1.setNameY("φ [rad]")

pp=P1D("particles")
for i in range(plist.size()):
      p=plist.get(i)
      eta=p.eta(); phi=p.phi()
      pp.add(eta,phi)
      
for k in range(jets.size()):
    jj=jets.get(k)
    cic= Circle(jj.eta(), jj.phi(), R)
    cic.setFill(1); cic.setColor(Color.red)
    c1.add(cic)
c1.draw(pp)