# View ProMC event in 3D
# HepSim script using DataMelt http://jwork.org/dmelt
# Part of =HepMC= : https://atlaswww.hep.anl.gov/hepsim/
# S.Chekanov (ANL)
from java.awt import Color,Font
from java.lang import *
from proto import FileMC            
from jhplot import HPlot3D,P2D
from hepsim import HepSim
from java.util import Random
from math import sqrt
import math,sys



# look at file 1, event 2
EventToView=2;
FileToView=1;

# magnetic field
# 1 Tesla = 10,000 Gauss 
# 1 kiloGauss = 1000 Gauss
field = 100;  # magnetic field in kilogauss
XYvolume=50;  # volume size in X or Y 
Zvolume=50;   # volume in Z

c1 = HPlot3D("=HepSim=",600,600)
c1.setGTitle("3D view of truth event")
c1.visible(True)
c1.setRange(-XYvolume,XYvolume,-XYvolume,XYvolume,-Zvolume,Zvolume)
c1.setNameX("X")
c1.setNameY("Y")
c1.setNameY("Z")

# step for helix
step    = 0.2;
vect=[0,0,0,0,0,0,0];
destep  = 0;
nmec    = 0;
ran=Random();
kX,kY,kZ,kPX,kPY,kPZ,kPP = 0,1,2,3,4,5,6

def helixStep(step,vect):
   vout=[0,0,0,0,0,0]
   h4    = field*2.99792e-4;
   rho   = -h4/vect[kPP] # p*charge 
   tet   = rho*step;
   tsint = tet*tet/6.0;
   sintt = 1.0 - tsint;
   sint  = tet*sintt;
   cos1t = tet/2.0;
   f1 = step*sintt;
   f2 = step*cos1t;
   f3 = step*tsint*vect[kPZ] # pZ 
   f4 = -tet*cos1t;
   f5 = sint;
   f6 = tet*cos1t*vect[kPZ];
   vout[kX]   = vect[kX]  + (f1*vect[kPX] - f2*vect[kPY]);
   vout[kY]   = vect[kY]  + (f1*vect[kPY] + f2*vect[kPX]);
   vout[kZ]   = vect[kZ]  + (f1*vect[kPZ] + f3);
   vout[kPX]  = vect[kPX] + (f4*vect[kPX] - f5*vect[kPY]);
   vout[kPY]  = vect[kPY] + (f4*vect[kPY] + f5*vect[kPX]);
   vout[kPZ]  = vect[kPZ] + (f4*vect[kPZ] + f6);
   return vout 


url="";  TotalEvents=1
default_www="http://mc.hep.anl.gov/asc/hepsim/events/pp/13tev/mg5_ttbar_bjet/"
if len(sys.argv)>1: 
   if sys.argv[1].startswith("http"): 
                            flist=HepSim.getList(sys.argv[1])
                            url=sys.argv[1]
   else: flist=FileList.get(sys.argv[1],"promc")
   if (len(sys.argv)>2): TotalEvents=int(sys.argv[2]) 
else:
   url=default_www; flist=HepSim.getList(url)
if len(flist)==0: print "Error: No input file!"; sys.exit(0)
else: print "Reading "+str(len(flist))+" files. Nr events= ",TotalEvents


print "Look at file=",FileToView
print "Look at event=",EventToView
if (FileToView>=len(flist)):
    print "File number is too large!. Max nummber=",Nfiles 
    sys.exit(0);

file=FileMC(url+flist[FileToView])         # get input file
header = file.getHeader()
chargeMap={}
# get charge and particle names
for j in range(header.getParticleDataCount()): 
     d = header.getParticleData(j) 
     pid = d.getId()
     mass = d.getMass()
     name = d.getName()
     chargeMap[pid]=int(d.getCharge()/3.0) 


un=float(header.getMomentumUnit()) # conversion units
lunit=float(header.getLengthUnit())
print "ProMC v=",file.getVersion(), "M unit=",un,"L unit=",lunit 
if (EventToView>=file.size()):
     print "Event number is too large! Max number=",file.size()
     sys.exit(0);
eve = file.read(EventToView)
pa = eve.getParticles()    # particle information

for j in range(pa.getPxCount()):
        hadrons=P2D("hadrons")
        hadrons.setSymbolSize(1)
        hadrons.setSymbolColor(Color.blue)
        leptons=P2D("leptons")
        leptons.setSymbolSize(1)
        leptons.setSymbolColor(Color.red)
        apdg=abs(pa.getPdgId(j))
        charge=chargeMap[pa.getPdgId(j)]
        if (apdg==12 or apdg==14 or apdg==16): continue # no neutrino
        if (pa.getStatus(j)==1 and charge !=0):
             print "processing particle=",j
             xL=pa.getX(j)/lunit
             yL=pa.getY(j)/lunit
             zL=pa.getZ(j)/lunit
             px=pa.getPx(j)/un
             py=pa.getPy(j)/un    
             pz=pa.getPz(j)/un 
             e=pa.getEnergy(j)/un
             p=math.sqrt(px*px+py*py+pz*pz)
             mass=pa.getMass(j)/un;
             vect[0]=xL
             vect[1]=yL
             vect[2]=zL
             vect[3]=px/p
             vect[4]=py/p
             vect[5]=pz/p
             vect[6]=p*charge
             getot=e
             gekin   = getot - mass;
             for k in range(400):
                  # build helix
                   vout=helixStep(step,vect); # make one helix step
                   destep = step* (0.0002*ran.nextGaussian()+0.00001);
                   gekin -= destep;
                   getot  = gekin + mass;
                   vect[6]= charge*sqrt(getot*getot - mass*mass);
                   vect[0]=vout[0]
                   vect[1]=vout[1]
                   vect[2]=vout[2]
                   vect[3]=vout[3]
                   vect[4]=vout[4]
                   vect[5]=vout[5]
                   if (abs(vect[0])<XYvolume):
                     if (abs(vect[1])<XYvolume):
                       if (abs(vect[2])<Zvolume):
                          if (apdg>100): hadrons.add(vect[0], vect[1], vect[2]) 
                          if (apdg<20):  leptons.add(vect[0], vect[1], vect[2])
        # draw objects
        c1.draw(hadrons)
        c1.draw(leptons)

stat = file.getStat()
cross=stat.getCrossSectionAccumulated()
erro=stat.getCrossSectionErrorAccumulated();
file.close()
print "Total cross section (pb)=",cross
#c1.update()


# c1.export(name+".svg");    print name+".svg created"
# sys.exit(0)