# Run this script using Jas4pp or DMelt Java programs
# This scripts makes 1000 events with single pions using:
# - Random -4.5<eta<4.5
# - Random phi
# - Random p in the range 0-10 GeV
#  
# Example using Jas4pp: "fpad particle_gun.py"
#
# May 2016. 
# S.Chekanov (ANL)

from java.io import BufferedOutputStream, BufferedReader, InputStreamReader, FileOutputStream
from java.net import *
from java.util import Random
from proto import *
from promc.io import *
from hepsim import *
from hephysics.particle import *
from java.util.zip import ZipEntry, ZipOutputStream
from java.lang import Math, Long, ClassLoader
import re

# output file
outFilename="outfile.promc"
momentum = 10 # random momentum up to 10 GeV
xPID = 211    # create pions
Ntot=1000     # number of events per file 

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()
# these units for energy and distance for varint convertion
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()
    # flat phi
    phi   = -Math.PI + 2*Math.PI*random.nextDouble()
    # flat Eta -4.5-4.5
    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)
    # do you need flat pt instead?
    # 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<T1: T2=A; T1=B
    R=4.0 # 4mm radius 
    Dx=T2*R*Math.cos(PI2*T1/T2)
    Dy=T2*R*Math.sin(PI2*T1/T2)
    Dt=0.0

    # add random single particle with fixed momentum
    addParticle(b_particles,unit,lunit,Id=4,M=G.mass(),M1=3,M2=3,PdgId=xPID,px=G.px(),py=G.py(),pz=G.pz(),e=G.e(),\
                status=1,weight=1,charge=xcharge,barcode=0,x=Dx,y=Dy,z=Dz,t=Dt,D1=0,D2=0)


    # fill this event with 4 particles, last is stable
    particles = b_particles.build()
    meta_event = event_builder.build()
    b_event.setEvent(meta_event)
    b_event.setParticles(particles)
    event = b_event.build()
    writeInfo(str(eventNo), event.toByteArray())

# file metadata
# description
writeInfo("promc_description",txtdescription.encode())


# write number of events
writeInfo("promc_nevents",str(Ntot).encode())

b_stat = ProMCStatFile.ProMCStat.newBuilder()
stat = b_stat.build()
#stat.setCrossSection(0)
#stat.setCrossSectionError(0)
writeInfo("statistics",stat.toByteArray())

s="Not set for this file"
# write Google templates
writeInfo("ProMCHeader.proto",s.encode())
writeInfo("ProMC.proto",s.encode())
writeInfo("ProMCStat.proto",s.encode())
writeInfo("ProMCDescription.proto",s.encode())
writeInfo("logfile.txt",s.encode())

# close 
zout.finish()
zout.close()
fout.close()

# open to view it
from probrowser import *
c=MainGui(outFilename)