<reponame>Algomorph/cvcalib
"""
Quaternion provides a class for manipulating quaternion objects.  This class provides:

    - a convenient constructor to convert to/from Euler Angles (RA,Dec,Roll) 
         to/from quaternions
    - class methods to multiply and divide quaternions 
"""

"""Copyright 2009 Smithsonian Astrophysical Observatory
    Released under New BSD / 3-Clause BSD License
    All rights reserved
"""
"""
    Modified 2012 by <NAME>
    Modified from version 0.3.1
    http://pypi.python.org/pypi/Quaternion/0.03
    Added _get_angle_axis to get the angle-axis representation
    Added _latlontoquat to get a rotation quat to ECEF from lat/lon
"""
"""
     Modified 2016 by <NAME>
     Added 
"""


        
import numpy as np
from math import cos, sin, radians, degrees, atan2, sqrt, acos, pi

class Quat(object):
    """
    Quaternion class
    
    Example usage::

     >>> from Quaternion import Quat
     >>> quat = Quat((12,45,45))
     >>> quat.ra, quat.dec, quat.roll
     (12, 45, 45)
     >>> quat.q
     array([ 0.38857298, -0.3146602 ,  0.23486498,  0.8335697 ])
     >>> q2 = Quat([ 0.38857298, -0.3146602 ,  0.23486498,  0.8335697])
     >>> q2.ra
     11.999999315925008


    Multiplication and division operators are overloaded for the class to 
    perform appropriate quaternion multiplication and division.

    Example usage::
    
     >>> q1 = Quat((20,30,40))
     >>> q2 = Quat((30,40,50))
     >>> q = q1 / q2

    Performs the operation as q1 * inverse q2

    Example usage::

     >>> q1 = Quat((20,30,40))
     >>> q2 = Quat((30,40,50))
     >>> q = q1 * q2


    :param attitude: initialization attitude for quat

    ``attitude`` may be:
      * another Quat
      * a 4 element array (expects x,y,z,w quat form)
      * a 3 element array (expects ra,dec,roll in degrees)
      * a 3x3 transform/rotation matrix

    """
    def __init__(self, attitude):
        self._q = None
        self._equatorial = None
        self._T = None
        # checks to see if we've been passed a Quat 
        if isinstance(attitude, Quat):
            self._set_q(attitude.q)
        else:
            # make it an array and check to see if it is a supported shape
            attitude = np.array(attitude)
            if len(attitude) == 4:
                self._set_q(attitude)
            elif attitude.shape == (3, 3):
                self._set_transform(attitude)
            elif attitude.shape == (3,):
                self._set_equatorial(attitude)
            elif attitude.shape == (2,):
                self._set_latlon(attitude)
            else:
                raise TypeError("attitude is not one of possible types (2, 3 or 4 elements, Quat, or 3x3 matrix)")
            

    def _set_q(self, q):
        """
        Set the value of the 4 element quaternion vector 

        :param q: list or array of normalized quaternion elements
        """
        q = np.array(q)
        if abs(np.sum(q ** 2) - 1.0) > 1e-6:
            raise ValueError('Quaternion must be normalized so sum(q**2) == 1; use Quaternion.normalize')
        self._q = (q if q[3] > 0 else -q)
        # Erase internal values of other representations
        self._equatorial = None
        self._T = None

    def _get_q(self):
        """
        Retrieve 4-vector of quaternion elements in [x, y, z, w] form
        
        :rtype: numpy array

        """
        if self._q is None:
            # Figure out q from available values, doing nothing others are not defined
            if self._equatorial is not None:
                self._q = self._equatorial2quat()
            elif self._T is not None:
                self._q = self._transform2quat()
        return self._q

    # use property to make this get/set automatic
    q = property(_get_q, _set_q)

    def _set_equatorial(self, equatorial):
        """Set the value of the 3 element equatorial coordinate list [RA,Dec,Roll]
            expects values in degrees
            bounds are not checked
        
            :param equatorial: list or array [ RA, Dec, Roll] in degrees
            
        """
        att = np.array(equatorial)
        ra, dec, roll = att#@UnusedVariable
        self._ra0 = ra
        if (ra > 180):
            self._ra0 = ra - 360
            self._roll0 = roll
        if (roll > 180):
            self._roll0 = roll - 360
        self._equatorial = att

    def _set_latlon(self, latlon):
        self._q = self._latlontoquat(latlon)

    def _get_equatorial(self):
        """Retrieve [RA, Dec, Roll]

        :rtype: numpy array
        """
        if self._equatorial is None:
            if self._q is not None:
                self._equatorial = self._quat2equatorial()
            elif self._T is not None:
                self._q = self._transform2quat()
                self._equatorial = self._quat2equatorial()
        return self._equatorial

    equatorial = property(_get_equatorial, _set_equatorial)

    def _get_ra(self):
        """Retrieve RA term from equatorial system in degrees"""
        return self.equatorial[0]

    def _get_dec(self):
        """Retrieve Dec term from equatorial system in degrees"""
        return self.equatorial[1]

    def _get_roll(self):
        """Retrieve Roll term from equatorial system in degrees"""
        return self.equatorial[2]

    ra = property(_get_ra)
    dec = property(_get_dec)
    roll = property(_get_roll)

    def _set_transform(self, T):
        """
        Set the value of the 3x3 rotation/transform matrix
        
        :param T: 3x3 array/numpy array
        """
        transform = np.array(T)
        self._T = transform

    def _get_transform(self):
        """
        Retrieve the value of the 3x3 rotation/transform matrix

        :returns: 3x3 rotation/transform matrix
        :rtype: numpy array
        
        """
        if self._T is None:
            if self._q is not None:
                self._T = self._quat2transform()
            elif self._equatorial is not None:
                self._T = self._equatorial2transform()
        return self._T

    transform = property(_get_transform, _set_transform)

    def _quat2equatorial(self):
        """
        Determine Right Ascension, Declination, and Roll for the object quaternion
        
        :returns: RA, Dec, Roll
        :rtype: numpy array [ra,dec,roll]
        """
        
        q = self.q
        q2 = self.q ** 2

        # # calculate direction cosine matrix elements from $quaternions
        xa = q2[0] - q2[1] - q2[2] + q2[3] 
        xb = 2 * (q[0] * q[1] + q[2] * q[3]) 
        xn = 2 * (q[0] * q[2] - q[1] * q[3]) 
        yn = 2 * (q[1] * q[2] + q[0] * q[3]) 
        zn = q2[3] + q2[2] - q2[0] - q2[1] 

        # #; calculate RA, Dec, Roll from cosine matrix elements
        ra = degrees(atan2(xb , xa)) ;
        dec = degrees(atan2(xn , sqrt(1 - xn ** 2)));
        roll = degrees(atan2(yn , zn)) ;
        if (ra < 0):
            ra += 360
        if (roll < 0):
            roll += 360

        return np.array([ra, dec, roll])


    def _quat2transform(self):
        """
        Transform a unit quaternion into its corresponding rotation matrix (to
        be applied on the right side).
        
        :returns: transform matrix
        :rtype: numpy array
        
        """
        x, y, z, w = self.q
        xx2 = 2 * x * x
        yy2 = 2 * y * y
        zz2 = 2 * z * z
        xy2 = 2 * x * y
        wz2 = 2 * w * z
        zx2 = 2 * z * x
        wy2 = 2 * w * y
        yz2 = 2 * y * z
        wx2 = 2 * w * x
        
        rmat = np.empty((3, 3), float)
        rmat[0, 0] = 1. - yy2 - zz2
        rmat[0, 1] = xy2 - wz2
        rmat[0, 2] = zx2 + wy2
        rmat[1, 0] = xy2 + wz2
        rmat[1, 1] = 1. - xx2 - zz2
        rmat[1, 2] = yz2 - wx2
        rmat[2, 0] = zx2 - wy2
        rmat[2, 1] = yz2 + wx2
        rmat[2, 2] = 1. - xx2 - yy2
        
        return rmat

    def _equatorial2quat(self):
        """Dummy method to return return quat. 

        :returns: quaternion
        :rtype: Quat
        
        """
        return self._transform2quat()
    
    def _equatorial2transform(self):
        """Construct the transform/rotation matrix from RA,Dec,Roll

        :returns: transform matrix
        :rtype: 3x3 numpy array

        """
        ra = radians(self._get_ra())
        dec = radians(self._get_dec())
        roll = radians(self._get_roll())
        ca = cos(ra)
        sa = sin(ra)
        cd = cos(dec)
        sd = sin(dec)
        cr = cos(roll)
        sr = sin(roll)
        
        # This is the transpose of the transformation matrix (related to translation
        # of original perl code
        rmat = np.array([[ca * cd, sa * cd, sd      ],
                              [-ca * sd * sr - sa * cr, -sa * sd * sr + ca * cr, cd * sr],
                              [-ca * sd * cr + sa * sr, -sa * sd * cr - ca * sr, cd * cr]])

        return rmat.transpose()

    def _transform2quat(self):
        """Construct quaternion from the transform/rotation matrix 

        :returns: quaternion formed from transform matrix
        :rtype: numpy array
        """

        # Code was copied from perl PDL code that uses backwards index ordering
        T = self.transform.transpose()  
        den = np.array([ 1.0 + T[0, 0] - T[1, 1] - T[2, 2],
                              1.0 - T[0, 0] + T[1, 1] - T[2, 2],
                              1.0 - T[0, 0] - T[1, 1] + T[2, 2],
                              1.0 + T[0, 0] + T[1, 1] + T[2, 2]])
        
        max_idx = np.flatnonzero(den == max(den))[0]

        q = np.zeros(4)
        q[max_idx] = 0.5 * sqrt(max(den))
        denom = 4.0 * q[max_idx]
        if (max_idx == 0):
            q[1] = (T[1, 0] + T[0, 1]) / denom 
            q[2] = (T[2, 0] + T[0, 2]) / denom 
            q[3] = -(T[2, 1] - T[1, 2]) / denom 
        if (max_idx == 1):
            q[0] = (T[1, 0] + T[0, 1]) / denom 
            q[2] = (T[2, 1] + T[1, 2]) / denom 
            q[3] = -(T[0, 2] - T[2, 0]) / denom 
        if (max_idx == 2):
            q[0] = (T[2, 0] + T[0, 2]) / denom 
            q[1] = (T[2, 1] + T[1, 2]) / denom 
            q[3] = -(T[1, 0] - T[0, 1]) / denom 
        if (max_idx == 3):
            q[0] = -(T[2, 1] - T[1, 2]) / denom 
            q[1] = -(T[0, 2] - T[2, 0]) / denom 
            q[2] = -(T[1, 0] - T[0, 1]) / denom 

        return q
        
    def _get_angle_axis(self):
        lim = 1e-12
        norm = np.linalg.norm(self.q)
        if norm < lim:
            angle = 0
            axis = [0, 0, 0]
        else:
            rnorm = 1.0 / norm
            angle = acos(max(-1, min(1, rnorm * self.q[3])));
            sangle = sin(angle)
            if sangle < lim:
                axis = [0, 0, 0]
            else:
                axis = (rnorm / sangle) * np.array(self.q[0:3])

            angle *= 2

        return (angle, axis)

    def _latlontoquat (self, latlon):
        q = np.zeros(4)
        
        lon = latlon[1] * (pi / 180.)
        lat = latlon[0] * (pi / 180.)
        zd2 = 0.5 * lon
        yd2 = -0.25 * pi - 0.5 * lat
        Szd2 = sin(zd2)
        Syd2 = sin(yd2)
        Czd2 = cos(zd2)
        Cyd2 = cos(yd2)
        q[0] = -Szd2 * Syd2
        q[1] = Czd2 * Syd2
        q[2] = Szd2 * Cyd2
        q[3] = Czd2 * Cyd2

        return q

    def __div__(self, quat2):
        """
        Divide one quaternion by another.
        
        Example usage::

         >>> q1 = Quat((20,30,40))
         >>> q2 = Quat((30,40,50))
         >>> q = q1 / q2

        Performs the operation as q1 * inverse q2

        :returns: product q1 * inverse q2
        :rtype: Quat

        """
        return self * quat2.inv()


    def __mul__(self, quat2):
        """
        Multiply quaternion by another.

        Example usage::

          >>> q1 = Quat((20,30,40))
          >>> q2 = Quat((30,40,50))
          >>> (q1 * q2).equatorial
          array([ 349.73395729,    76.25393056,  127.61636727])

        :returns: product q1 * q2
        :rtype: Quat

        """
        q1 = self.q
        q2 = quat2.q
        mult = np.zeros(4)
        mult[0] = q1[3] * q2[0] - q1[2] * q2[1] + q1[1] * q2[2] + q1[0] * q2[3]
        mult[1] = q1[2] * q2[0] + q1[3] * q2[1] - q1[0] * q2[2] + q1[1] * q2[3]
        mult[2] = -q1[1] * q2[0] + q1[0] * q2[1] + q1[3] * q2[2] + q1[2] * q2[3]
        mult[3] = -q1[0] * q2[0] - q1[1] * q2[1] - q1[2] * q2[2] + q1[3] * q2[3]
        return Quat(mult)

    def inv(self):
        """
        Invert the quaternion 

        :returns: inverted quaternion
        :rtype: Quat
        """
        return Quat([self.q[0], self.q[1], self.q[2], -self.q[3]])
    
    @staticmethod
    def slerp(qa, qb, t):
        # Calculate angle between them.
        #qa.w * qb.w + qa.x * qb.x + qa.y * qb.y + qa.z * qb.z;
        cosHalfTheta = np.dot(qa.q,qb.q)
        #if qa=qb or qa=-qb then theta = 0 and we can return qa
        if (abs(cosHalfTheta) >= 1.0):
            return Quat(np.copy(qa.q));
        #Calculate temporary values.
        halfTheta = acos(cosHalfTheta);
        sinHalfTheta = sqrt(1.0 - cosHalfTheta*cosHalfTheta);
        
        #if theta = 180 degrees then result is not fully defined
        #we could rotate around any axis normal to qa or qb
        if(abs(sinHalfTheta) < 0.001):
            return Quat(qa.q * 0.5 + qb.q * 0.5);
        
        ratioA = sin((1 - t) * halfTheta) / sinHalfTheta;
        ratioB = sin(t * halfTheta) / sinHalfTheta; 
        #calculate Quaternion for general case.
        return Quat(qa.q * ratioA + qb.q * ratioB);
    
    @staticmethod
    def normalize(array):
        """ 
        Normalize a 4 element array/list/numpy.array for use as a quaternion
        
        :param quat_array: 4 element list/array
        :returns: normalized array
        :rtype: numpy array
    
        """
        quat = np.array(array)
        return quat / np.sqrt(np.dot(quat, quat))
