#!/usr/bin/env python3
#
# filter_bhgen.py - script to pass over GENBEAM 'BHgen' output events
#                   from hdgeant to select pairs and triplet events for
#                   full simulation by hdgeant4.
#
# author: richard.t.jones at uconn.edu
# version: july 3, 2024

import hddm_s
import sys
import os
import math

deg = math.pi/180
minpolarangle = 0.75 * deg
mintantheta = math.tan(minpolarangle)

def usage():
   print("usage: filter_bhgen.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:
   discard = 0
   for pev in rec.getPhysicsEvents():
      pev.deleteHitViews()
   for rea in rec.getReactions():
      if rea.weight == 0:
         continue
      rea.deleteVertices(1)
      for vtx in rea.getVertices():
         for ori in vtx.getOrigins():
            ori.t = 0
            ori.vx = 0
            ori.vy = 0
            ori.vz = 0
            pro = vtx.getProducts()
            for p in range(2):
               for mom in pro[p].getMomenta():
                  tantheta = (mom.px**2 + mom.py**2)**0.5 / (mom.pz + 1e-99)
                  if tantheta < mintantheta:
                     discard += 1
         if discard == 0:
            hostream.write(rec)
            count += 1

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