# 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"""
	def __init__(self, dims, pos):
		# Note: Later updating of any of these will NOT affect the rest!
		# i.e. if redefine dims, then nr and far will not change
		# accordingly. May want to include a 'resize' method to overcome
		# this
		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. Includes shared attribute 'model_limits' which
	pushes the boundary of the model when a Detector is defined"""
	# The cuboid which contains all the detectors: the model limits
	model_limits = Box([0,0,0], [0,0,0])
	def __init__(self, dims, pos):
		# Run Box constructor
		Box.__init__(self, dims, pos)
		# Adjust model boundaries to accomodate detector
		for i in range(len(self.far[:])):
			if self.far[i] < Detector.model_limits.far[i]: 
				Detector.model_limits.far[i] = self.far[i] 
		for i in range(len(self.nr[:])):
			if self.nr[i] > Detector.model_limits.nr[i]:
				Detector.model_limits.nr[i] = self.nr[i]
		# Define an array to store all particle hits
		self.hits = {}

class Particle:
	"""Defines a particle

	Initialises a particle with a position and velocity vector"""
	def __init__(self, pos, vel):
		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"""
	def __init__(self, vol):
		# 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