#!/usr/bin/env python3
#
# filter_postcol.py - script to pass over GENBEAM 'postcol' output events
#                     from hdgeant to assign circular polarization to the
#                     beam photons with the appropriate sign: + if the first
#                     argument is odd, or - if the first argument is even.
#
# author: richard.t.jones at uconn.edu
# version: july 3, 2024

import hddm_s
import sys
import os

zplane = 65 # cm, project the beam photon downstream to this plane
clight = 29.9792 # cm / ns

def usage():
   print("usage: filter_postcol.py <configNo> <infile.hddm> <outfile.hddm>")
   print(" where <configNo> is a positive integer, right circular polarization")
   print(" for <configNo>=1,2 and left for <configNo>=3,4.")
   sys.exit(1)

try:
   if int(sys.argv[1]) < 3:
      beampol = +1
   else:
      beampol = -1
   infile = sys.argv[2]
   outfile = sys.argv[3]
   histream = hddm_s.istream(infile)
   hostream = hddm_s.ostream(outfile)
except:
   usage()

count = 0
for rec in histream:
   for rea in rec.getReactions():
      for bea in rea.getBeams():
         for pol in bea.getPolarizations():
            pol.Px = 0
            pol.Py = 0
            pol.Pz = beampol
      rea.deleteVertices(1)
      for vtx in rea.getVertices():
         for ori in vtx.getOrigins():
            for pro in vtx.getProducts():
               for mom in pro.getMomenta():
                  ori.t += (zplane - ori.vz) / clight
                  ori.vx += mom.px / mom.pz * (zplane - ori.vz)
                  ori.vy += mom.py / mom.pz * (zplane - ori.vz)
                  ori.vz = zplane
                  mom.E = (mom.px**2 + mom.py**2 + mom.pz**2)**0.5
                  mass2 = mom.E**2 - mom.px**2 - mom.py**2 - mom.pz**2
                  if mass2 > 1e-12:
                     pev = rec.getPhysicsEvents()
                     print(f"strange photon mass {mass2**0.5} seen in event {pev[0].eventNo}")
                     print(mom.E, mom.px, mom.py, mom.pz)
                     print(mom.E**2 - mom.pz**2, mom.px**2, mom.py**2, mass2)
               for pol in pro.getPolarizations():
                  pol.Px = 0
                  pol.Py = 0
                  pol.Pz = beampol
         hostream.write(rec)
         count += 1

print(f"filter_postcol.py transferred {count} events from {infile} to {outfile}")