#! /usr/bin/env python

## a program that does an example 2D simulation in which particles
## (gamma rays) come out of a point source and bounce against various
## objects in the geometry; one of the objects is a detector.

## the output can be viewed in many ways; here are some:
## python sim-first.py > sim-first.out
## grep HIT: sim-first.out > hits
## grep STEP: sim-first.out > steps
## gnuplot
## gnuplot> plot 'steps' using 2:3 with dots
## gnuplot> replot 'hits' using 2:3 with points
## gnuplot> splot 'steps' using 2:3:4 with dots
## gnuplot> replot 'hits' using 2:3:4 with points

import math
import random

n_particles = 1000
object_list = []

## physical constants
c_light_mks = 299792458.0
h_planck_mks = (6.626068E-34)
#oneEVinJoules = (1.60217646E-19); /* conversion factor */
m_electron_mks = (9.10938188E-31)

## particles start at (0, -10) and the ones we are interested in go in
## the +Z direction
source_x = 0
source_z = -10
## specify the geometry of the world
world_box = [[-20, -20], [20, 20]]

def __main__():
  ## when we want it reproducible we force the seed; comment this line
  ## out to seed on current time
  random.seed(31415)
  prepare_geometry()
  for i in range(n_particles):
    run_particle()


## the geometries for now are simple -- just a few rectangles
def prepare_geometry():
  ## each object is a list of corner (x, z) coordinates
  object_list.append([[-2, -1], [3, 3], 80, 'Pb'])
  object_list.append([[4, 2], [7, 2.5], 50, 'Fe'])

## run a particle through the simulation -- the first approach is to
## do a simple bouncing instead of compton scattering
def run_particle():
  ## for now we collimate at an agle of -30 to 30 degrees in the Z
  ## direction, which means the mean angle should be at 90 degrees
  angle_deg = random.uniform(90-30, 90+30)
  vx = c_light_mks * math.cos(angle_deg*math.pi/180.0)
  vz = c_light_mks * math.sin(angle_deg*math.pi/180.0)
  print 'START:', angle_deg, vx/c_light_mks, vz/c_light_mks
  track_particle(source_x, source_z, vx, vz)


## simple tracking of a particle by taking a bunch of steps at time dt
def track_particle(x0, z0, vx, vz):
  time = 0
  dt = 4.0e-10
  E = get_initial_energy()
  x = source_x
  z = source_z
  done = 0                          # done is 1 when we exit the world
  while not done:
    (done, x, z, vx, vz, E, time) = get_next_hit(x, z, vx, vz, E, time, dt)
    print_hit_info(x, z, E, time, dt)



## see when we get the next hit -- take steps until we hit a box or
## exit the world.  return 0 of hit is a box, 1 if it's exiting the
## world
def get_next_hit(x0, z0, vx0, vz0, E0, t0, dt):
  x = x0
  z = z0
  vx = vx0
  vz = vz0
  E = E0
  t = t0
  while point_is_in_box(x, z, world_box):
    t += dt
    x += vx*dt
    z += vz*dt
    print_step_info(x, z, E, t, dt)
    for geometry_object in object_list:
      if point_is_in_box(x, z, geometry_object):
        ## should we do compton scattering?  depends on the electron
        ## density of the material
        if we_should_do_compton(geometry_object):
          (vx, vz, E) = do_compton(vx, vz, E)
          if E < 0:                       # E < 0 is similar to "out of world"
            print 'ABSORBED:', x, z, vx, vz, E, t
            return (1, x, z, vx, vz, E, t)
          else:
            return (0, x, z, vx, vz, E, t)        # ball is still in play
  ## if we exit this loop it means we have exited the world
  return (1, x, z, vx, vz, E, t)


## return an initial energy; there are many possible spectra you could
## give
def get_initial_energy():
  ## I could implement various spectra; for now I just do the positron
  ## annihilation
  E_kev = 511                           # simple positron annihilation
  E_mks = 1.609e-19 * E_kev * 1000.0;   # convert kev to Joules
  return E_mks


def point_is_in_box(x, z, box):       # check if point is in a box
##   print 'POINT_IS_IN_BOX:', x, z, box,
  if x > box[0][0] and x < box[1][0] and z > box[0][1] and z < box[1][1]:
##     print '   ---   1'
    return 1
##   print '   ---   0'
  return 0

def do_compton(vx, vz, E):
  v = math.sqrt(vx * vx + vz * vz)
  original_angle = math.atan2(vz/v, vx/v)
  scat_angle = choose_compton_theta(E)

  new_angle = original_angle + scat_angle
  vx = v*math.cos(new_angle)
  vz = v*math.sin(new_angle)
  delta_lambda = h_planck_mks/(m_electron_mks*c_light_mks*c_light_mks)*(1-math.cos(scat_angle))
  E = 1.0 / (delta_lambda/h_planck_mks + 1.0/(E))
  return (vx, vz, E)

## figure out what theta should be used for the compton scattering
def choose_compton_theta(E):
  theta_min = 0.0
  theta_max = math.pi
  x = random.uniform(theta_min, theta_max)
  y = random.uniform(theta_dist_klein_nishina(E, theta_min),
                     theta_dist_klein_nishina(E, theta_max))
  while y > theta_dist_klein_nishina(E, x):
    x = random.uniform(theta_min, theta_max)
    y = random.uniform(theta_dist_klein_nishina(E, theta_min),
                       theta_dist_klein_nishina(E, theta_max))
  return x;


## use the klein-nishina formula for the 
def theta_dist_klein_nishina(E, theta):
  R_E = 2.817940E-10	# electron radius in meters
  ## first define some convenience functions
  def alpha(E):
    return E/(m_electron_mks*c_light_mks*c_light_mks)
  def oneminuscos(theta):
    return 1 - math.cos(theta)
  def oneminuscos_sq(theta):
    omct = 1 - math.cos(theta);
    return omct*omct;

  basic_den = (1+alpha(E)*oneminuscos(theta)) # denominator
  ae = alpha(E)
  return math.cos(theta)*(R_E*R_E/2.0) \
    * 1.0/(basic_den*basic_den) \
    * (1.0 + math.cos(theta)*math.cos(theta) + ae*ae*oneminuscos_sq(theta)/basic_den);


## using the electron density, decide if we should do compton
## scattering
def we_should_do_compton(geometry_object):
  r = random.uniform(0, 100)
  if r < geometry_object[2]:
    return 1
  else:
    return 0

## print current info
def print_hit_info(x, z, E, time, dt):
  print 'HIT: %12.6g   %12.6g   %12.6g     %12.6g' % (x, z, time, E)

def print_step_info(x, z, E, time, dt):
  print 'STEP: %12.6g   %12.6g   %12.6g     %12.6g' % (x, z, time, E)


__main__()