#########################################################################
#
# Script: basic_kinematics.py
#
# Author: Igor Volobouev, Feb 2005
#
# Purpose: implements formulae related to relativistic kinematics,
#          spherical rotation interpolation, and other basic math
#
#########################################################################

"""Module basic_kinematics -- classes for representing vectors in 3 and 4
dimensions and performing basic operations on them such as rotations,
boosts, scalar and vector products, etc. Also contains functions for
spherical rotation interpolation."""

import math
import cmath

class V3:
    "Basic 3-vector."
    def __init__(self, x=0.0, y=0.0, z=0.0):
        self.x = x*1.0
        self.y = y*1.0
        self.z = z*1.0

    def norm2(self):
        "Spatial length of the vector, squared."
        return self.x*self.x + self.y*self.y + self.z*self.z

    def norm(self):
        "Spatial length of the vector."
        return math.sqrt(self.norm2())

    def direction(self):
        "Direction of the vector as a 3-vector with unit length."
        if (self.__nonzero__()):
            return self / self.norm()
        else:
            return V3(1.0, 0.0, 0.0)

    def eta(self):
        "Pseudo-rapidity."
        if (self.__nonzero__()):
            return cmath.atanh(self.z/self.norm()).real
        else:
            return 0.0

    def phi(self):
        "Azimuthal angle."
        return math.atan2(self.y, self.x)

    def cosTheta(self):
        "Cosine of the polar angle."
        if (self.__nonzero__()):
            return self.z/self.norm()
        else:
            return 0.0

    def theta(self):
        "Polar angle."
        cosTheta = self.cosTheta()
        if (math.fabs(cosTheta) < 0.99):
            return math.acos(cosTheta)
        else:
            # acos would loose too much numerical precision
            th = math.asin(math.sqrt((self.x*self.x + \
                                      self.y*self.y)/self.norm2()))
            if (cosTheta > 0.0):
                return th
            else:
                return math.pi - th

    def dec(self):
        "Declination."
        return 90.0 - self.theta()*180.0/math.pi

    def ra(self):
        "Right ascension in degrees."
        ang = self.phi()*180.0/math.pi
        if (ang < 0.0):
            ang += 360.0
        return ang

    def angle(self, other):
        "Angle between two vectors."
        u = self.direction()
        v = other.direction()
        cosa = sprod3(u, v)
        if (math.fabs(cosa) < 0.99):
            return math.acos(cosa)
        else:
            # acos would loose too much numerical precision
            if (cosa > 0.0):
                return 2.0*math.asin(abs(v - u)/2.0)
            else:
                return math.pi - 2.0*math.asin(abs(-v - u)/2.0)

    def __repr__(self):
        return '{' + str(self.x) + ', ' + str(self.y) + \
               ', ' + str(self.z) + '}'

    def __nonzero__(self):
        return int(self.x != 0.0 or self.y != 0.0 or self.z != 0.0)

    def __neg__(self):
        return V3(-self.x, -self.y, -self.z)

    def __pos__(self):
        return V3(self.x, self.y, self.z)

    def __abs__(self):
        return self.norm()

    def __eq__(self, other):
        return self.x == other.x and \
               self.y == other.y and \
               self.z == other.z

    def __ne__(self, other):
        return not (self == other)

    def __add__(self, other):
        return V3(self.x+other.x, self.y+other.y, self.z+other.z)

    def __sub__(self, other):
        return V3(self.x-other.x, self.y-other.y, self.z-other.z)

    def __mul__(self, other):
        return V3(self.x*other, self.y*other, self.z*other)

    def __rmul__(self, other):
        return self*other

    def __div__(self, other):
        if (other == 0.0):
            raise ZeroDivisionError, "3-vector divided by zero"
        return V3(self.x*1.0/other, self.y*1.0/other, self.z*1.0/other)

    # Mutators
    def __iadd__(self, other):
        self.x += other.x
        self.y += other.y
        self.z += other.z
        return self

    def __imul__(self, other):
        self.x *= other
        self.y *= other
        self.z *= other
        return self

    def __isub__(self, other):
        self.x -= other.x
        self.y -= other.y
        self.z -= other.z
        return self

    def __idiv__(self, other):
        if (other == 0.0):
            raise ZeroDivisionError, "3-vector divided by zero"
        self.x /= (other * 1.0)
        self.y /= (other * 1.0)
        self.z /= (other * 1.0)
        return self


class FourMomentum:
    "Time-like 4-vector."
    def __init__(self, *args):
        arglen = len(args)
        if (arglen == 0):
            # 0 arguments: default constructor
            self.p = V3()
            self.m = 0.0
            self.esign = 1.0
        elif (arglen == 2):
            # 2 arguments: either (e, p) or (p, m)
            arg0 = args[0]
            arg1 = args[1]
            if isinstance(arg0, V3):
                # called using (p, m)
                self.p = arg0
                self.m = arg1*1.0
                self.esign = 1.0
            elif isinstance(arg1, V3):
                # called using (e, p)
                self.p = arg1
                self.m = math.sqrt(arg0*arg0 - arg1.norm2())
                if (arg0 >= 0.0):
                    self.esign = 1.0
                else:
                    self.esign = -1.0
            else:
                raise TypeError, "invalid argument type"
        elif (arglen == 3):
            # called using (p, m, sign)
            self.p = args[0]
            self.m = args[1]*1.0
            if (args[2] >= 0.0):
                self.esign = 1.0
            else:
                self.esign = -1.0
        else:
            raise SyntaxError, "wrong # of args"
        if (self.m < 0.0):
            raise ValueError, "mass can't be negative"

    def _getEnergy(self):
        return self.esign * math.sqrt(self.p.norm2() + self.m*self.m)
    e = property(_getEnergy, doc="Energy component of a time-like 4-vector")

    def beta(self):
        "Particle speed in units of c."
        return self.p.norm()/self.e

    def gamma(self):
        "Particle relativistic gamma."
        return self.e*1.0/self.m

    def betagamma(self):
        "Product of particle speed in units of c and relativistic gamma."
        return self.p.norm()/self.m

    def __repr__(self):
        assert self.esign != 0.0
        if (self.esign > 0.0):
            csign = '+'
        else:
            csign = '-'
        return '{' + csign + ', ' + str(self.p) + ', ' + str(self.m) + '}'

    def __nonzero__(self):
        return int(bool(self.p) or self.m != 0.0)

    def __pos__(self):
        return FourMomentum(self.p, self.m, self.esign)

    def __neg__(self):
        # Note that we keep the mass positive
        return FourMomentum(-self.p, self.m, -self.esign)

    def __abs__(self):
        return self.m

    def __eq__(self, other):
        if (self or other):
            return self.p == other.p and self.m == other.m and \
                   self.esign == other.esign
        else:
            return 1

    def __ne__(self, other):
        return not (self == other)

    def __add__(self, other):
        return FourMomentum(self.e+other.e, self.p+other.p)

    def __sub__(self, other):
        return FourMomentum(self.e-other.e, self.p-other.p)

    def __mul__(self, other):
        return FourMomentum(self.p*other, self.m*abs(other), self.esign*other)

    def __rmul__(self, other):
        return self*other

    def __div__(self, other):
        if (other == 0.0):
            raise ZeroDivisionError, "4-vector divided by zero"
        return FourMomentum(self.p*1.0/other, self.m*1.0/abs(other),
                            self.esign*1.0/other)

    # Mutators
    def __iadd__(self, other):
        newe = self.e + other.e
        self.m = invmass(self, other)
        self.p += other.p
        if (newe >= 0.0):
            self.esign = 1.0
        else:
            self.esign = -1.0
        return self

    def __isub__(self, other):
        newe = self.e - other.e
        msq = self.m*self.m + other.m*other.m - 2.0*sprod4(self,other)
        # The following will raise a math domain
        # error if msq is less than zero
        self.m = math.sqrt(msq)
        self.p -= other.p
        if (newe >= 0.0):
            self.esign = 1.0
        else:
            self.esign = -1.0
        return self

    def __imul__(self, other):
        newe = self.esign*other
        self.m *= abs(other)
        self.p *= other
        if (newe >= 0.0):
            self.esign = 1.0
        else:
            self.esign = -1.0
        return self

    def __idiv__(self, other):
        if (other == 0.0):
            raise ZeroDivisionError, "4-vector divided by zero"
        newe = self.esign*1.0/other
        self.m /= (1.0*abs(other))
        self.p /= (1.0*other)
        if (newe >= 0.0):
            self.esign = 1.0
        else:
            self.esign = -1.0
        return self


class Quaternion:
    "Helpfull class for representing rotations in 3d."
    def __init__(self, s, vec):
        self.s = s
        self.v = vec

    def conjugate(self):
        "Returns conjugated quaternion."
        return Quaternion(self.s, -self.v)

    def norm2(self):
        "Norm of a quaternion squared."
        return self.v.norm2() + self.s*self.s

    def norm(self):
        "Quaternion norm."
        return math.sqrt(self.norm2())

    def __nonzero__(self):
        return int(self.s != 0.0 or bool(self.v))

    def __eq__(self, other):
        return self.s == other.s and self.v == other.v

    def __ne__(self, other):
        return not (self == other)

    def __abs__(self):
        return self.norm()

    def __add__(self, other):
        return Quaternion(self.s + other.s, self.v + other.v)

    def __sub__(self, other):
        return Quaternion(self.s - other.s, self.v - other.v)

    def __mul__(self, other):
        if isinstance(other, Quaternion):
            return Quaternion(self.s*other.s - sprod3(self.v, other.v), \
                   other.v*self.s + self.v*other.s + vprod3(self.v, other.v))
        else:
            return Quaternion(self.s*other, self.v*other)

    def __repr__(self):
        return '{' + str(self.s) + ', ' + str(self.v) + '}'

    def __pos__(self):
        return Quaternion(self.s, self.v)

    def __neg__(self):
        return Quaternion(-self.s, -self.v)

    def __div__(self, other):
        if isinstance(other, Quaternion):
            return self * other.__inverse__()
        else:
            # "other" must be a scalar
            return Quaternion(self.s*1.0/other, self.v/other)

    def __inverse__(self):
        return self.conjugate() / self.norm2()


class Rotation:
    """Representation of rotations in 3d.

0 constructor arguments: identity rotation
1 argument : either a Quaternion or Rotation object
2 arguments: rotation axis and angle
4 arguments: right ascension and declination of x and z axes

Use () operator to rotate a 3-vector.
Multiply to perform sequential rotations: (r1*r2)(v) = r1(r2(v)).
Unary ~ operator can be used to find an inverse rotation."""
    # Basically, this class just makes sure that the quaternion
    # used to represent the rotation stays normalized. It also
    # can create the normalized quaternion from the right ascension
    # and declination coordinates used throughout the LAT software.
    def __init__(self, *args):
        arglen = len(args)
        if (arglen == 0):
            self.quat = Quaternion(1.0, V3(0.0, 0.0, 0.0))
        
        elif (arglen == 1):
            arg0 = args[0]
            if isinstance(arg0, Quaternion):
                # Normalize the quaternion
                self.quat = arg0 / arg0.norm()
            elif isinstance(arg0, Rotation):
                # Quaternion is already normalized
                self.quat = arg0.quat
            else:
                raise TypeError, "invalid argument type"
        
        elif (arglen == 2):
            if (args[0]):
                axis      = args[0].direction()
                halftheta = args[1] / 2.0
                self.quat = Quaternion(math.cos(halftheta), \
                                       axis*math.sin(halftheta))
            else:
                raise ValueError, "null rotation axis"
        
        elif (arglen == 4):
            xdir = dirFromRaDec(args[0], args[1])
            zdir = dirFromRaDec(args[2], args[3])
            # We assume that the z direction is more
            # precise. This is a rather arbitrary assumption,
            # of course. Then we renormalize the axes.
            ydir = vprod3(zdir, xdir).direction()
            xdir = vprod3(ydir, zdir)
            # Convert rotation matrix (xdir, ydir, zdir) into
            # the quaternion representation
            t = 1.0 + xdir.x + ydir.y + zdir.z
            if (t > 0.01):
                # Standard formula result will be numerically stable
                x = ydir.z - zdir.y
                y = zdir.x - xdir.z
                z = xdir.y - ydir.x
                w = t
            else:
                # It is better to use alternative representations
                if (xdir.x >= ydir.y and xdir.x >= zdir.z):
                    x = 1.0 + xdir.x - ydir.y - zdir.z
                    y = ydir.x + xdir.y
                    z = zdir.x + xdir.z
                    w = ydir.z - zdir.y
                elif (ydir.y >= xdir.x and ydir.y >= zdir.z):
                    x = ydir.x + xdir.y
                    y = 1.0 + ydir.y - xdir.x - zdir.z
                    z = zdir.y + ydir.z
                    w = zdir.x - xdir.z
                else:
                    x = zdir.x + xdir.z
                    y = zdir.y + ydir.z
                    z = 1.0 + zdir.z - xdir.x - ydir.y
                    w = xdir.y - ydir.x
            q = Quaternion(w, V3(x, y, z))
            self.quat = q / q.norm()
        
        else:
            raise SyntaxError, "wrong # of args"

    def axis(self):
        "Rotation axis as a 3-vector with a unit length."
        return self.quat.v.direction()

    def angle(self):
        "Rotation angle in radians."
        return 2.0*math.atan2(self.quat.v.norm(), self.quat.s)

    def __call__(self, vec):
        # Rotation acts on a 3 or 4-vector (or their subclasses)
        vcopy = +vec
        if (self):
            try:
                # 4-vector branch
                p = vec.p
                is4 = True
            except AttributeError:
                # 3-vector branch
                p = vec
                is4 = False
            rotated = ((self.quat*Quaternion(0.0,p))*self.quat.conjugate()).v
            if is4:
                vcopy.p = rotated
            else:
                vcopy = rotated
        return vcopy

    def __nonzero__(self):
        return int(bool(self.quat.v))

    def __mul__(self, other):
        # Two successful rotations. Associative but not commutative.
        return Rotation(self.quat * other.quat)

    def __pos__(self):
        return Rotation(self)

    def __invert__(self):
        "Calculates the inverse rotation"
        return Rotation(self.quat.conjugate())

    def __repr__(self):
        return str(self.quat)

    def __eq__(self, other):
        return self.quat == other.quat or self.quat == -other.quat

    def __ne__(self, other):
        return not (self == other)


class LorentzBoost:
    """Representation of Lorentz boosts.

1 constructor argument: FourMomentum. Boost into the rest system of the
                        given particle.
2 arguments: boost direction (3-vector not necessarily of unit length)
             and rapidity.

Use () operator to boost a FourMomentum.
Unary ~ operator can be used to find an inverse boost."""
    def __init__(self, *args):
        arglen = len(args)
        if (arglen == 0):
            # Default constructor
            self.direction = V3(0.0, 0.0, 0.0).direction()
            self.rapidity = 0.0
        elif (arglen == 1):
            # Constructor using a 4-vector which defines the rest system
            fourMom = args[0]
            if (fourMom.m == 0.0):
                raise ValueError, \
                      "can't boost into the rest system of a massless particle"
            self.direction = fourMom.p.direction()
            self.rapidity = cmath.atanh(fourMom.p.norm()/fourMom.e).real;
        elif (arglen == 2):
            # Constructor using direction and rapidity
            self.direction = args[0].direction()
            self.rapidity = args[1]
        else:
            raise SyntaxError, "wrong # of args"

    def __repr__(self):
        return '<' + str(self.direction) + ', ' + str(self.rapidity) + '>'

    def __nonzero__(self):
        return int(self.rapidity != 0.0)

    def __eq__(self, other):
        if (self or other):
            return (self.direction == other.direction and \
                    self.rapidity  == other.rapidity) or \
                   (self.direction == -other.direction and \
                    self.rapidity  == -other.rapidity)
        else:
            return 1

    def __ne__(self, other):
        return not (self == other)

    def __pos__(self):
        return LorentzBoost(self.direction, self.rapidity)

    def __invert__(self):
        "Calculates the inverse boost"
        return LorentzBoost(self.direction, -self.rapidity)

    def __call__(self, fourMom):
        "Boosts the given particle"
        # Boost acts on a 4-vector or its subclass. Changes
        # momentum only, not mass or sign of energy.
        pcopy = +fourMom
        if (self):
            # Momentum component parallel to the boost
            par  = sprod3(self.direction, fourMom.p)
            # Momentum perpendicular to the boost
            perp = fourMom.p - self.direction*par
            # Boosted parallel component
            parPrime = math.cosh(self.rapidity)*par - \
                       fourMom.e*math.sinh(self.rapidity)
            # Boosted momentum
            pcopy.p = perp + self.direction*parPrime
        return pcopy


def dirFromRaDec(ra, dec):
    "Makes a unit vector given its right ascension in degrees and declination."
    theta = (90.0 - dec)*math.pi/180.0
    phi   = ra*math.pi/180.0
    sinTheta = math.sin(theta)
    return V3(sinTheta*math.cos(phi), sinTheta*math.sin(phi), math.cos(theta))

def dirFromLatLon(lat, lon):
    "Makes a unit vector given its latitude in degrees and longitude."
    return dirFromRaDec(lon, lat)

def sprod3(p1, p2):
    "Scalar product of 3-vectors, unit metric."
    return p1.x*p2.x + p1.y*p2.y + p1.z*p2.z

def vprod3(p1, p2):
    "Vector product of 3-vectors, unit metric."
    return V3(p1.y*p2.z-p1.z*p2.y, p1.z*p2.x-p1.x*p2.z, p1.x*p2.y-p1.y*p2.x)

def sprod4(p1, p2):
    "Scalar product of 4-vectors."
    return p1.e*p2.e - sprod3(p1.p, p2.p)

def qprod4(q1, q2):
    "Dot product of two quaternions."
    return q1.s*q2.s + sprod3(q1.v, q2.v)

def invmass(p1, p2, *args):
    "Invariant mass of two or more 4-vectors."
    arglen = len(args)
    if (arglen == 0):
        return math.sqrt(p1.m*p1.m + p2.m*p2.m + 2.0*sprod4(p1, p2))
    else:
        p3 = args[0]
        if (arglen == 1):
            return math.sqrt(p1.m*p1.m + p2.m*p2.m + p3.m*p3.m + \
                             2.0*sprod4(p1, p2) + 2.0*sprod4(p1, p3) + \
                             2.0*sprod4(p2, p3))
        else:
            p4 = args[1]
            if (arglen == 2):
                return math.sqrt(p1.m*p1.m + p2.m*p2.m + \
                                 p3.m*p3.m + p4.m*p4.m + \
                                 2.0*sprod4(p1, p2) + 2.0*sprod4(p1, p3) + \
                                 2.0*sprod4(p2, p3) + 2.0*sprod4(p1, p4) + \
                                 2.0*sprod4(p2, p4) + 2.0*sprod4(p3, p4))
            else:
                # At this point it is probably easier to add 4-vectors
                # since the number of sprod4 operations is O(N^2)...
                sum = p1 + p2
                for p in args:
                    sum += p
                return sum.m

def sinx_over_x(x):
    "Numerically stable implementation of sin(x)/x."
    if (math.fabs(x) < 1.0e-8):
        return 1.0;
    else:
        return math.sin(x)/x;

def slerp(r1, r2, t):
    """Spherical linear rotation interpolation.

Arguments are rotations at t = 0, 1, and time."""
    # This is a fast, precise, and numerically stable
    # implementation of the Ken Shoemake's quaternion
    # interpolation formula (Proceedings of SIGGRAPH 85,
    # pp 245-254).
    q1 = r1.quat
    cosa = qprod4(q1, r2.quat)
    if (cosa >= 0.0):
        q2 = r2.quat
    else:
        q2 = -r2.quat
        cosa = -cosa
    if (cosa < 0.99):
        a = math.acos(cosa)
    else:
        a = 2.0*math.asin(abs(q1 - q2)/2.0)
    onemt = 1.0 - t
    c0 = sinx_over_x(a)
    c1 = onemt*sinx_over_x(a*onemt)/c0
    c2 = t*sinx_over_x(a*t)/c0
    return Rotation(q1*c1 + q2*c2)

def slerp2(r1, r2, t):
    "Spherical linear rotation interpolation. Useful for testing slerp."
    # This function has a lot cleaner code than "slerp",
    # and really shows you what is going on. Unfortunately,
    # it is also about three times slower than "slerp".
    # Both functions produce identical results.
    if (qprod4(r1.quat, r2.quat) >= 0.0):
        rot = (~r1)*r2
    else:
        rot = (~r1)*Rotation(-r2.quat)
    return r1*Rotation(rot.axis(), t*rot.angle())

def scubic(r1, r2, r3, r4, t):
    """Spherical cubic rotation interpolation.

Arguments are rotations at t = 0, 1/3, 2/3, 1, and time."""
    # Note that this formula is different from the popular "squad"
    # formula of Shoemake. Shoemake's formula is a Bezier curve.
    # This one passes through the control points.
    return slerp(slerp(r1, r4, t), slerp(r2, r3, 3.0*t-1.0), 4.5*t*(1.0-t))