# ParticleModelClasses.py
# ---------------------------------------------------------------------
# Defines the following classes, useful for creating models of particle
# detectors. See each classes docstrign for more information.
# 
# Vector    - A list of numbers that behave mathematically as a vector
# Box       - A cuboid in the model, Detector inherits from this
# Detector  - A cuboid detector in the model
# Particle  - A particle with velocity and position variables
# CosmicRay - A particle created in a random position in a space with
#             velocity according to the azimuthal Cos^2 distribution 

# IMPORT SOME NECESSARY FUNCTIONS

from math import *
from random import random

# CLASSES

class Vector(list):
	"""Mathematical vector of any size

	Changes behaviour of lists so as to behave more like a vector w.r.t
	addition etc. Methods are described seperately. Builds on code by A.
	Pletzer in the ActiveState Python Cookbook"""
	def __add__(self, other):
		"Returns an element by element sum in a Vector"
		return Vector([i+j for i,j in zip(self, other)])
	def __neg__(self):
		"Returns Vector with all elements negatived"
		return Vector([-i for i in self])
	def __sub__(self, other):
		"Returns Vector of the difference of two vectors"
		return Vector([i-j for i, j in zip(self,other)])
	def __mul__(self, other):
		"Returns element by element multiplication in a Vector"
		try: return Vector([i*j for i,j in zip(self,other)])
		# Must be a scalar
		except: return Vector([i*other for i in self])
	def __rmul__(self, other):
		"Since multiplcation commutes, __rmul__() does as __mul__()"
		return self*other
	def __div__(self, other):
		"Divides, element by element"
		try: return Vector([i/j for i,j in zip(self, other)])
		# Must be scalar
		except: return Vector([i/other for i in self])
	def magnitude(self):
		"Returns magnitude of the vector"
		return sqrt(sum([x**2 for x in self]))
	def unit(self):
		"Returns unit vector in direction" 
		return Vector(self / self.magnitude)
	def dot(self, other):
		"Returns dot product of this and another vector"
		return sum(self * other)
	def cross(self, other):
		"""Returns the cross product of this and another vector.
		 Currently only works for Vectors with 3 values"""
		if len(other) != 3 or len(self) != 3:
			raise ValueError, 'cross method currently only works for 3 valued Vectors'
		else:
			return Vector(self[1]*other[2]-self[2]*other[1], \
					self[0]*other[2]-self[2]*other[0], \
					self[0]*other[1]-self[1]*other[0])
		

class Box:
	"""Defines a cuboid in 3D

	First argument when a Box is created is a 3 valued vector giving the
	dimensions of the cuboid, the second argument gives the position of the
	centre. If no arguments given a unit box centred on origin is created"""
	def __init__(self, dims=[1,1,1], pos=[0,0,0]):
		# Note: Do not change the dimensions and position directly,
		# since the the corner positions will not be updated and vise
		# versa. Use the 'set' methods below 
		self.set_dims_pos(dims,pos)
	def set_corners(self, nr, far):
		"Use to redefine corners of box"
		self.nr = Vector(nr)
		self.far = Vector(far)
		self.dims = Vector([n-f for n,f in zip(nr,far)])
		self.pos = Vector([0.5*d+f for d,f in zip(self.dims,far)])
	def set_dims_pos(self, dims, pos):
		"Use to redefine position and dimensions of box"
		self.dims = Vector(dims)
		self.pos = Vector(pos)
		self.nr = Vector([0.5*d+p for d,p in zip(dims,pos)])
		self.far = Vector([-0.5*d+p for d,p in zip(dims,pos)])
	def inside(self, pnt):
		"Returns True if a point is inside box, False otherwise"
		if pnt[0] > self.nr[0] or pnt[0] < self.far[0]: return False
		if pnt[1] > self.nr[1] or pnt[1] < self.far[1]: return False
		if pnt[2] > self.nr[2] or pnt[2] < self.far[2]: return False
		return True

class Detector(Box):
	"""Defines a cuboid shaped detector in 3D

	Inheriting from the Box class, the dimensions and position are
	defined as for a Box."""
	def __init__(self, dims, pos):
		# Run Box constructor
		Box.__init__(self, dims, pos)
		# Define an array to store all particle hits
		self.hits = {}

class Particle:
	"""Defines a particle

	Initialises a particle with a position and velocity vector. If no
	arguments given a stationary particle at the origin is created"""
	def __init__(self, pos=[0,0,0], vel=[0,0,0]):
		self.pos = Vector(pos)
		self.vel = Vector(vel)

class CosmicRay(Particle):
	"""Defines a cosmic ray particle

	 Creates a particle in a cuboid with random direction and
	displacement and traces back to point of entry. Parameter vol is
	instance of Box or Detector class. If no argument is passed, CosmicRay
	is created at origin"""
	def __init__(self, vol=Box([0,0,0],[0,0,0])):
		# Azimuthal angle according to Cos^2 probability distribution.
		# This uses a Von Neumann acceptance, rejection algorithm, see
		# Section 33.3 Phys. Rev. 2004
		while 1:
			self.phi = random() * pi / 2
			y = random()
			if y <= cos(self.phi)**2: break
		self.theta = random() * 2.0 * pi # Rotational angle of particle
		# Calculate unit velocity vector
		self.vel = Vector([ sin(self.phi) * sin(self.theta), \
				sin(self.phi) * cos(self.theta), \
				-cos(self.phi) ])
		# Place particle in random place in vol
		self.pos = Vector([ random() * vol.dims[0] + vol.far[0], \
				random() * vol.dims[1] + vol.far[1], \
				random() * vol.dims[2] + vol.far[2] ])
		# Traceback the particle position until it is just inside of the
		# vol
		self.traceback(vol)
	# Define traceback function which places particle at entry point to
	# vol
	def traceback(self, vol):
		# Start by skipping 10 mm at a time
		while 1:
			self.pos = self.pos - 10 * self.vel
			if not vol.inside(self.pos): break
		# Place particle back inside detectors vol
		while not vol.inside(self.pos):
			self.pos = self.pos + self.vel