# Fourvec class: implementation of four-vectors
#
# J. A. Templon, NIKHEF-K, Amsterdam
#
# All masses in MeV/c^2, momenta in MeV/c, energies in MeV, and
#    angles in degrees
#
# Usage: initialize with
#
# from fourvec import Fourvec
#
#    your_fourvec = Fourvec(energy,type1,vec,type2)
#
# 'type1' specifies how the particle mass is defined:
#     type1 = 'rest'  -- 'energy' given is the particle rest mass
#     type1 = 'total' -- 'energy' given is the particle total energy
# 'type2' specifies how the momentum vector is defined:
#     type2 = 'polar' -- 'vec' given is the list
#             vec = [ Pmag, thet, phi ]
#     type2 = 'cart'  -- 'vec' given is the list
#             vec = [ px, py, pz ]
#
#    Pmag is the absolute magnitude of the 3-momentum (in MeV/c)
#    thet and phi are the azimuthal angles of the momentum vector (degrees)
#
# attributes:
#
# fourvec.E - the total energy of the particle in MeV
# fourvec.px - the x component of the momentum (in MeV/c)
# fourvec.py - the y component of the momentum (in MeV/c)
# fourvec.pz - the z component of the momentum (in MeV/c)
#
# methods modifying attributes:
#
#   These methods directly assign momentum 3-vec to existing
#   fourvec object; rest mass is preserved, TOTAL ENERGY RECOMPUTED!!
#
# fourvec.assn_p_cart(vec) - vec is list [px, py, pz]
# fourvec.assn_p_pola(vec) - vec is list [Pmag, thet, phi]
#
# methods returning output:
#
# fourvec.m0()    - returns the particle rest mass in MeV/c^2
# fourvec.T()     - returns the particle kinetic energy in MeV
# fourvec.P()     - returns the absolute magnitude of the 3-momentum (MeV/c)
# fourvec.thet()  - returns the polar angle of the momentum 3-vector (degrees)
# fourvec.phi()   - returns the azim. angle of the momentum 3-vector (degrees)
# fourvec.pvec()  - returns the momentum 3- vector as a list [ px, py, pz ]
# fourvec.gamma() - returns the Lorentz factor
# fourvec.beta()  - returns the reduced velocity
#
# operators:
#
# "-" (__sub__)  - subtraction of two four-vectors
# "+" (__add__)  - addition of two four-vectors

import math, string

DEGRAD = 180.0/math.pi     # conversion radians -> degrees

class Fourvec:

    def __init__(self,energy,type1,vec,type2) :

#-> initialize the momentum vector and magnitude

	ptype = string.lower(type2)
	if ptype == 'cart' :
	    self.px = vec[0]
	    self.py = vec[1]
	    self.pz = vec[2]
	    Pmag = self.px*self.px + self.py*self.py + self.pz*self.pz
	    Pmag = math.sqrt(Pmag)
	elif ptype == 'polar' :
	    thetarad = vec[1] / DEGRAD
	    phirad   = vec[2] / DEGRAD
	    Pmag     = vec[0]
	    self.px = Pmag * math.sin(thetarad) * math.cos(phirad)
	    self.py = Pmag * math.sin(thetarad) * math.sin(phirad)
	    self.pz = Pmag * math.cos(thetarad)
	else :
	    raise ValueError, \
		  'Error in Fourvec init: undefined 3-vec type ' \
		  + '\"' + ptype + '\"'

#-> now compute or set the total energy

	etype = string.lower(type1)
	if etype == 'rest' :
	    self.E = math.sqrt(pow(energy,2) + pow(Pmag,2))
	elif etype == 'total' :
	    self.E = energy
	else :
	    raise ValueError, \
		  'Error in Fourvec init: undefined energy type ' \
		  + '\"' + etype + '\"'
# end of __init__

    def pvec(self) :    # returns the momentum 3-vector

	return [ self.px, self.py, self.pz ]

    def P(self) :       # returns the magnitude of the 3-momentum

	p_mag = self.px*self.px + self.py*self.py + self.pz*self.pz
	p_mag = math.sqrt(p_mag)
	return p_mag

    def m0(self) :      # returns the rest mass

	rm = math.sqrt(pow(self.E,2) - pow(self.P(),2))
	return rm

    def T(self) :

	ekin = self.E - self.m0()
	return ekin

    def gamma(self) :    # return the Lorentz factor

	return self.E / self.m0()

    def beta(self) :     # return the reduced velocity

	invgam2 = 1.0 / ( pow(self.gamma(),2) )
	return math.sqrt( 1.0 - invgam2 )

    def thet(self) :
	return DEGRAD * math.acos( self.pz / self.P() )

    def phi(self) :
	if self.px == 0 and self.py == 0 :
	    return 0.0
	else :
	    return DEGRAD * math.atan2( self.py, self.px )

    def assn_p_cart(self,vec) :   # setup 4v with cartesian p input

	rm = self.m0()
	self.px = vec[0]
	self.py = vec[1]
	self.pz = vec[2]
	self.E = math.sqrt(pow(rm,2) + pow(self.P(),2))

    def assn_p_pola(self,Pmag,thet,phi) : # setup 4v with polar p input

	rm = self.m0()
	thetarad = thet / DEGRAD
	phirad   = phi  / DEGRAD
	self.E = math.sqrt( pow(rm,2) + pow(Pmag,2) )
	self.px = Pmag * math.sin(thetarad) * math.cos(phirad)
	self.py = Pmag * math.sin(thetarad) * math.sin(phirad)
	self.pz = Pmag * math.cos(thetarad)

    def __add__(self, x) :
	rE = self.E + x.E
	
	rpx = self.px + x.px
	rpy = self.py + x.py
	rpz = self.pz + x.pz
	rv  = Fourvec(rE,'total',[rpx,rpy,rpz],'cart')
	return rv

    def __sub__(self, x) :
	rE = self.E - x.E
	rpx = self.px - x.px
	rpy = self.py - x.py
	rpz = self.pz - x.pz
	rv  = Fourvec(rE,'total',[rpx,rpy,rpz],'cart')
	return rv


def dot(x,y) :
    total = x[0]*y[0] + x[1]*y[1] + x[2]*y[2]
    return total

# END of Fourvec class