# Source code for rotmat

    '''
    Routines for working with rotation matrices
    '''

    '''
    comment
    author :  Thomas Haslwanter
    date :    June-2017
    '''
    __version__ = '2.0'

    import numpy as np
    import sympy
    from collections import namedtuple

    # The following construct is required since I want to run the module as a script
    # inside the skinematics-directory
    import os
    import sys
    sys.path.append( os.path.join( os.path.dirname(__file__), os.path.pardir ) )

    from skinematics import quat

    # For deprecation warnings
    import deprecation
    import warnings
    #warnings.simplefilter('always', DeprecationWarning)

    def R(axis=0, angle=90) :
        '''Rotation matrix for rotation about a cardinal axis.
        The argument is entered in degree.

        Parameters
        ----------
        axis : skalar
                Axis of rotation, has to be 0, 1, or 2
        alpha : float
                rotation angle [deg]

        Returns
        -------
        R : rotation matrix, for rotation about the specified axis

        Examples
        --------
        >>> rotmat.R(axis=0, alpha=45)
        array([[ 1.        ,  0.        ,  0.        ],
               [ 0.        ,  0.70710678, -0.70710678],
               [ 0.        ,  0.70710678,  0.70710678]])

        >>> rotmat.R(axis=0)
        array([[ 1.        ,  0.        ,  0.        ],
               [ 0.        ,  0.        , -1.        ],
               [ 0.        ,  1.        ,  0.       ]])

        >>> rotmat.R(1, 45)
        array([[ 0.70710678,  0.        ,  0.70710678],
               [ 0.        ,  1.        ,  0.        ],
               [-0.70710678,  0.        ,  0.70710678]])

        >>> rotmat.R(axis=2, alpha=45)
        array([[ 0.70710678, -0.70710678,  0.        ],
               [ 0.70710678,  0.70710678,  0.        ],
               [ 0.        ,  0.        ,  1.        ]])
        '''

        # convert from degrees into radian:
        a_rad = np.deg2rad(angle)

        if axis == 0:
            R = np.array([[1,             0,            0],
                          [0, np.cos(a_rad), -np.sin(a_rad)],
                          [0, np.sin(a_rad),  np.cos(a_rad)]])

        elif axis == 1:
            R = np.array([[ np.cos(a_rad), 0, np.sin(a_rad) ],
                          [            0,  1,             0 ],
                          [-np.sin(a_rad), 0, np.cos(a_rad) ]])

        elif axis == 2:
            R = np.array([[np.cos(a_rad), -np.sin(a_rad), 0],
                          [np.sin(a_rad),  np.cos(a_rad), 0],
                          [            0,             0,  1]])

        else:
            raise IOError('"axis" has to be 0, 1, or 2')
        return R

    def R_s(axis=0, angle='alpha'):
        '''
        Symbolic rotation matrix about the given axis, by an angle with the given name

        Returns
        -------
            R_symbolic : symbolic matrix for rotation about the given axis

        Examples
        --------

        >>> R_yaw = R_s(axis=2, angle='theta')

        >>> R_aero = R_s(2) * R_s(1) * R_s(0)

        '''

        alpha = sympy.Symbol(angle)

        if axis == 0:
            R_s =  sympy.Matrix([[1,                0,                 0],
                             [0, sympy.cos(alpha), -sympy.sin(alpha)],
                             [0, sympy.sin(alpha),  sympy.cos(alpha)]])

        elif axis == 1:
            R_s = sympy.Matrix([[sympy.cos(alpha),0, sympy.sin(alpha)],
                             [0,1,0],
                             [-sympy.sin(alpha), 0, sympy.cos(alpha)]])

        elif axis == 2:
            R_s = sympy.Matrix([[sympy.cos(alpha), -sympy.sin(alpha), 0],
                             [sympy.sin(alpha), sympy.cos(alpha), 0],
                             [0, 0, 1]])

        else:
            raise IOError('"axis" has to be 0, 1, or 2')
        return R_s

    def sequence(R, to ='Euler'):
        '''
        This function takes a rotation matrix, and calculates
        the corresponding angles for sequential rotations.

        R_Euler = R3(gamma) * R1(beta) * R3(alpha)

        Parameters
        ----------
        R : ndarray, 3x3
            rotation matrix
        to : string
            Has to be one the following:
            - Euler ... Rz * Rx * Rz
            - Fick ... Rz * Ry * Rx
            - aero ... same as "Fick"
            - Helmholtz ... Ry * Rz * Rx

        Returns
        -------
        alpha : first rotation(right-most matrix)
        beta : second rotation
        gamma : third rotation (left-most matrix)

        Notes
        -----
        The following formulas are used:

        Euler:

        .. math::
            \\beta   = -arcsin(\\sqrt{ R_{13}^2 + R_{23}^2 }) * sign(R_{23})

            \\gamma = arcsin(\\frac{R_{13}}{sin\\beta})

            \\alpha   = arcsin(\\frac{R_{31}}{sin\\beta})

        aero / Fick:

        .. math::

            \\theta   = arctan(\\frac{R_{21}}{R_{11}})

            \\phi = arcsin(R_{31})

            \\psi   = arctan(\\frac{R_{32}}{R_{33}})

        Note that it is assumed that psi < pi !

        Helmholtz:

        .. math::
            \\theta = arcsin(R_{21})

            \\phi = -arcsin(\\frac{R_{31}}{cos\\theta})

            \\psi = -arcsin(\\frac{R_{23}}{cos\\theta})

        Note that it is assumed that psi < pi

        '''

        if to=='Fick' or to=='aero':
            gamma = np.arctan2(R[1,0], R[0,0])
            alpha = np.arctan2(R[2,1], R[2,2])
            beta = -np.arcsin(R[2,0])

        elif to == 'Helmholtz':
            gamma = np.arcsin(R[1,0])
            beta = -np.arcsin(R[2,0]/np.cos(gamma))
            alpha = -np.arcsin(R[1,2]/np.cos(gamma))

        elif to == 'Euler':
            epsilon = 1e-6
            beta = - np.arcsin(np.sqrt(R[0,2]**2 + R[1,2]**2)) * np.sign(R[1,2])
            if beta < 1e-6:
                beta = 0
                gamma = 0
                alpha = np.arcsin(R[1,0])
            else:
                gamma = np.arcsin(R[0,2]/np.sin(beta))
                alpha = np.arcsin(R[2,0]/np.cos(beta))

        else:
            print('\nSorry, only know: \naero, \nFick, \nHelmholtz, \nEuler.\n')
            raise IOError

        # Return the parameter-angles
        Angles = namedtuple('Angles', ['first', 'second', 'third'])
        return Angles(alpha, beta, gamma)

    # ------------ Everything below here is deprecated -------------------------------
    @deprecation.deprecated(deprecated_in="2.0", removed_in="2.2",
                            current_version=__version__,
                            details="Use the ``R(0, angle)`` function instead")
    def R1(psi):
        '''Rotation about the 1-axis.
        Deprecated, use R(0, psi)'''
        return R(axis=0, angle=psi)

    @deprecation.deprecated(deprecated_in="2.0", removed_in="2.2",
                            current_version=__version__,
                            details="Use the ``R(1, angle)`` function instead")
    def R2(phi):
        '''Rotation about the 2-axis.
        Deprecated, use R(1, phi)'''
        return R(axis=1, angle=phi)

    @deprecation.deprecated(deprecated_in="2.0", removed_in="2.2",
                            current_version=__version__,
                            details="Use the ``R(2, angle)`` function instead")
    def R3(theta):
        '''Rotation about the 3-axis.
        Deprecated, use R(2, theta)'''

    @deprecation.deprecated(deprecated_in="2.0", removed_in="2.2",
                            current_version=__version__,
                            details="Use the ``sequence(R, 'Euler')`` function instead")
    def rotmat2Euler(R):
        '''
        This function takes a rotation matrix, and calculates the corresponding Euler-angles.
        Deprecated. Use "sequence(R, 'Euler')" '''

        return sequence(R, 'Euler')

    @deprecation.deprecated(deprecated_in="2.0", removed_in="2.2",
                            current_version=__version__,
                            details="Use the ``sequence(R, 'Fick')`` function instead")
    def rotmat2Fick(R):
        '''
        This function takes a rotation matrix, and calculates the corresponding Fick-angles.
        Deprecated. Use "sequence(R, 'Fick')" '''

        return sequence(R, 'Fick')

    @deprecation.deprecated(deprecated_in="2.0", removed_in="2.2",
                            current_version=__version__,
                            details="Use the ``sequence(R, 'Helmholtz')`` function instead")
    def rotmat2Helmholtz(R):
        '''
        This function takes a rotation matrix, and calculates the corresponding Helmholtz-angles.
        Deprecated. Use "sequence(R, 'Helmholtz')" '''
        return sequence(R, 'Helmholtz')

    @deprecation.deprecated(deprecated_in="2.0", removed_in="2.2",
                            current_version=__version__,
                            details="Use the ``R_s(0, 'psi')`` function instead")
    def R1_s():
        '''
        Symbolic rotation matrix about the 1-axis, by an angle psi
        Deprecated. Use R_s(0, 'psi')'''
        return R_s(0, 'psi')

    @deprecation.deprecated(deprecated_in="2.0", removed_in="2.2",
                            current_version=__version__,
                            details="Use the ``R_s(1, 'phi')`` function instead")
    def R2_s():
        '''
        Symbolic rotation matrix about the 2-axis, by an angle phi
        Deprecated. Use R_s(1, 'phi')'''
        return R_s(1, 'phi')

    @deprecation.deprecated(deprecated_in="2.0", removed_in="2.2",
                            current_version=__version__,
                            details="Use the ``R_s(2, 'alpha')`` function instead")
    def R3_s():
        '''
        Symbolic rotation matrix about the 3-axis, by an angle theta
        Deprecated. Use R_s(2, 'theta')'''
        return R_s(2, 'theta')

    def convert(rMat, to ='quat'):
        '''
        Converts a rotation matrix to the corresponding quaternion.
        Assumes that R has the shape (3,3), or the matrix elements in columns

        Parameters
        ----------
        rMat : array, shape (3,3) or (N,9)
            single rotation matrix, or matrix with rotation-matrix elements.
        to : string
            Currently, only 'quat' is supported

        Returns
        -------
        outQuat : array, shape (4,) or (N,4)
            corresponding quaternion vector(s)

        Notes
        -----

        .. math::
             \\vec q = 0.5*copysign\\left( {\\begin{array}{*{20}{c}}
            {\\sqrt {1 + {R_{11}} - {R_{22}} - {R_{33}}} ,}\\\\
            {\\sqrt {1 - {R_{11}} + {R_{22}} - {R_{33}}} ,}\\\\
            {\\sqrt {1 - {R_{11}} - {R_{22}} + {R_{33}}} ,}
            \\end{array}\\begin{array}{*{20}{c}}
            {{R_{32}} - {R_{23}}}\\\\
            {{R_{13}} - {R_{31}}}\\\\
            {{R_{21}} - {R_{12}}}
            \\end{array}} \\right)

        More info under
        http://en.wikipedia.org/wiki/Quaternion

        Examples
        --------

        >>> rotMat = array([[cos(alpha), -sin(alpha), 0],
        >>>    [sin(alpha), cos(alpha), 0],
        >>>    [0, 0, 1]])
        >>> rotmat.convert(rotMat, 'quat')
        array([[ 0.99500417,  0.        ,  0.        ,  0.09983342]])

        '''

        if to != 'quat':
            raise IOError('Only know "quat"!')

        if rMat.shape == (3,3) or rMat.shape == (9,):
            rMat=np.atleast_2d(rMat.ravel()).T
        else:
            rMat = rMat.T
        q = np.zeros((4, rMat.shape[1]))

        R11 = rMat[0]
        R12 = rMat[1]
        R13 = rMat[2]
        R21 = rMat[3]
        R22 = rMat[4]
        R23 = rMat[5]
        R31 = rMat[6]
        R32 = rMat[7]
        R33 = rMat[8]

        q[1] = 0.5 * np.copysign(np.sqrt(1+R11-R22-R33), R32-R23)
        q[2] = 0.5 * np.copysign(np.sqrt(1-R11+R22-R33), R13-R31)
        q[3] = 0.5 * np.copysign(np.sqrt(1-R11-R22+R33), R21-R12)
        q[0] = np.sqrt(1-(q[1]**2+q[2]**2+q[3]**2))

        return q.T

    if __name__ == '__main__':
        testmat = np.array([[np.sqrt(2)/2, -np.sqrt(2)/2, 0],
                       [np.sqrt(2)/2,  np.sqrt(2)/2, 0],
                       [0, 0, 1]])
        Fick = rotmat2Fick(testmat)
        print(Fick)

        angles = sequence(testmat, to='aero')
        print(angles)

        print('Done testing')
        '''
        correct = np.r_[[0,0,np.pi/4]]
        print(R())
        print(R1(45))
        from pprint import pprint
        pprint(R_s(axis=0))
        pprint(R_s(axis=1, angle='phi'))

        print(R(1,45))
        print(R_s(1, 'gamma'))

        R2(30)
        print(R1(40))
        a = np.r_[np.cos(0.1), 0,0,np.sin(0.1)]
        print('The inverse of {0} is {1}'.format(a, quat.q_inv(a)))

        '''