# Single particle gun implemented using Jython/Python and Java
# Run this script using DMelt Java program
# S.Chekanov (ANL)

from java.io import *
from java.util import *
from java.net import *
from java.util import *
from proto import *
from promc.io import *
from hepsim import *
from java.util import *
from hephysics.particle import *
from java.util.zip import *
from java.lang import *
import re

outFilename="outfile.promc"
momentum = 1 # 1 GeV
xPID = 211   # pions
Ntot=100     # number of events

print "output=",outFilename

# write a single entry
def writeInfo(key,data):
      entry = ZipEntry(key)
      zout.putNextEntry(entry)
      entry.setSize(len(data))
      zout.write(data)
      zout.closeEntry()

# add one particle
def addParticle(b,unit,lunit,Id,M,M1,M2,PdgId,px,py,pz,e,status,weight,charge,barcode,x,y,z,t,D1,D2):
         b.addId(Id)
         b.addMass( (int)(M*unit) )
         b.addMother1(M1)
         b.addMother2(M2)
         b.addPdgId(PdgId)
         b.addPx( (int)(px*unit) )
         b.addPy( (int)(py*unit) )
         b.addPz( (int)(pz*unit) )
         b.addEnergy( (int)(e*unit) )
         b.addStatus(status)
         b.addWeight(weight)
         b.addCharge(charge)
         b.addBarcode(barcode)
         b.addX(  (int)(x*lunit) )
         b.addY(  (int)(y*lunit) )
         b.addZ(  (int)(z*lunit) )
         b.addT(  (int)(t*lunit) )
         b.addDaughter1(D1)
         b.addDaughter2(D2)



fout = FileOutputStream(outFilename)
zout = ZipOutputStream(BufferedOutputStream(fout))

random=Random()
PI2=2*Math.PI
PI=Math.PI
deg= 180 / Math.PI
txtdescription="Single particles. PID="+str(xPID)+" max E(GeV)="+str(momentum) 
print txtdescription

# create a header
b_desc= ProMCDescriptionFile.ProMCDescription.newBuilder()
b_desc.setVersion(4L)
b_desc.setEvents(Ntot)
b_desc.setTimestamp(1L)
b_desc.setDescription(txtdescription)
writeInfo("version",Long.toString(4).encode())

b_header = ProMCHeaderFile.ProMCHeader.newBuilder()
unit = 1000*100
lunit = 1000
b_header.setId1(xPID)
b_header.setId2(xPID)
b_header.setCrossSection(1)
b_header.setCrossSectionError(0)
b_header.setLengthUnit(lunit)
b_header.setMomentumUnit(unit)
b_header.setName("Single particles")
b_header.setNameLengthUnit("mm")
b_header.setNameMomentumUnit("GeV")


# get PDG table from the jar file (not external)
# Add PDG table to the output file
loader = ClassLoader.getSystemClassLoader()
stream = loader.getResourceAsStream("probrowser/resources/particle.tbl")
reader = BufferedReader(InputStreamReader(stream))
line = reader.readLine()
xmass=0
xcharge=0
while line is not None:
        line = reader.readLine()
        line=line.strip()
        parts=re.findall(r'\S+', line)
        if (len(parts)!=6): continue
        # print parts[0]," ",parts[1]," ",parts[2]," ",parts[3]," ",parts[4]," ",parts[5]
        b_pdg=ProMCHeaderFile.ProMCHeader.ParticleData.newBuilder()
        b_pdg.setId(int(parts[0]))
        b_pdg.setName(parts[1].strip())
        b_pdg.setMass(float(parts[3]))
        b_pdg.setWidth(float(parts[4]))
        b_pdg.setLifetime(float(parts[5]))
        charge=int(parts[2])
        b_pdg.setCharge( charge  )
        pdg=b_pdg.build()
        b_header.addParticleData(pdg) # add to the header file  
        if (xPID==int(parts[0])): 
                xmass=float(parts[3])
                xcharge= charge
                break
reader.close()

# build header now
header = b_header.build()
writeInfo("header",header.toByteArray())
desc = b_desc.build()
writeInfo("description",desc.toByteArray())

# make collision events:
oldPercentComplete = 0
for  eventNo in range(Ntot):
    percentComplete = (eventNo * 100.0 / Ntot)
    if(percentComplete % 5 == 0 and percentComplete > oldPercentComplete): 
         print percentComplete,"% complete"
         oldPercentComplete = percentComplete
    b_event = ProMC.ProMCEvent.newBuilder()
    event_builder=ProMC.ProMCEvent.Event.newBuilder()
    b_particles = ProMC.ProMCEvent.Particles.newBuilder()
    count=1

    addParticle(b_particles,unit,lunit,Id=0,M=0,M1=0,M2=0,PdgId=90,px=0,py=0,pz=0,e=14000,\
                     status=11,weight=1,charge=0,barcode=0,x=0,y=0,z=0,t=0,D1=0,D2=0)

    # proton -> 
    addParticle(b_particles,unit,lunit,Id=1,M=0.938272,M1=0,M2=0,PdgId=2212,px=0,py=0,pz=7000,e=7000,\
                     status=4,weight=1,charge=1,barcode=0,x=0,y=0,z=0,t=0,D1=0,D2=0)

    # proton <- 
    addParticle(b_particles,unit,lunit,Id=2,M=0.938272,M1=0,M2=0,PdgId=2212,px=0,py=0,pz=-7000,e=7000,\
                     status=4,weight=1,charge=1,barcode=0,x=0,y=0,z=0,t=0,D1=0,D2=0)

    # some gluon
    addParticle(b_particles,unit,lunit,Id=3,M=0,M1=0,M2=0,PdgId=21,px=0,py=0,pz=0,e=1400,\
                     status=4,weight=1,charge=0,barcode=0,x=0,y=0,z=0,t=0,D1=0,D2=0)

          
    # add single final-state particle with random Phi, Eta
    G=LParticle()
    phi   = -Math.PI + 2*Math.PI*random.nextDouble()
    eta   = -4.5 + 9*random.nextDouble()
    theta=2.0*Math.atan(Math.exp(-1*eta))
    G.setMass(xmass)
    G.setCharge(xcharge)
    G.setThetaPhiP(theta,phi,momentum)
    # G.setPtEtaPhiM(pt, eta, phi, xmass) 

    # Change origin too. mm in Z (-150 - 150 mm) 
    # assume a particle was produced uniformly a cylinder with R=4 mm
    Dz   = -150.0+300.0*random.nextDouble()
    A = random.nextDouble()
    B = random.nextDouble()
    T1=A
    T2=B
    if T2