'''
Tools for Manipulating and Converting 3D Rotations

By Omid Alemi
Created: June 12, 2017

Adapted from that matlab file...
'''

import math
import numpy as np

def deg2rad(x):
    return x/180*math.pi


def rad2deg(x):
    return x/math.pi*180


def rotmat2quat(R):
    """
    Converts a rotation matrix to a quaternion
    Matlab port to python for evaluation purposes
    https://github.com/asheshjain399/RNNexp/blob/srnn/structural_rnn/CRFProblems/H3.6m/mhmublv/Motion/rotmat2quat.m#L4
    Args
    R: 3x3 rotation matrix
    Returns
    q: 1x4 quaternion
    """
    rotdiff = R - R.T;

    r = np.zeros(3)
    r[0] = -rotdiff[1,2]
    r[1] =  rotdiff[0,2]
    r[2] = -rotdiff[0,1]
    sintheta = np.linalg.norm(r) / 2;
    r0 = np.divide(r, np.linalg.norm(r) + np.finfo(np.float32).eps );

    costheta = (np.trace(R)-1) / 2;

    theta = np.arctan2( sintheta, costheta );

    q      = np.zeros(4)
    q[0]   = np.cos(theta/2)
    q[1:] = r0*np.sin(theta/2)

    q_norm = np.linalg.norm(q);

    return q/q_norm

def quat2expmap(q):
    """
    Converts a quaternion to an exponential map
    Matlab port to python for evaluation purposes
    https://github.com/asheshjain399/RNNexp/blob/srnn/structural_rnn/CRFProblems/H3.6m/mhmublv/Motion/quat2expmap.m#L1
    Args
    q: 1x4 quaternion
    Returns
    r: 1x3 exponential map
    Raises
    ValueError if the l2 norm of the quaternion is not close to 1
    """
    if (np.abs(np.linalg.norm(q)-1)>1e-3):
        raise(ValueError, "quat2expmap: input quaternion is not norm 1")

    sinhalftheta = np.linalg.norm(q[1:])
    coshalftheta = q[0]

    r0    = np.divide( q[1:], (np.linalg.norm(q[1:]) + np.finfo(np.float32).eps));
    theta = 2 * np.arctan2( sinhalftheta, coshalftheta )
    theta = np.mod( theta + 2*np.pi, 2*np.pi )

    if theta > np.pi:
        theta =  2 * np.pi - theta
        r0  = -r0

    r = r0 * theta
    return r


class Rotation():
    def __init__(self,rot, param_type, **params):
        self.rotmat = []
        if param_type == 'euler':
            self._from_euler(rot[0],rot[1],rot[2], params)
        elif param_type == 'expmap':
            self._from_expmap(rot[0], rot[1], rot[2], params)

    def _from_euler(self, alpha, beta, gamma, params):
        '''Expecting degress'''

        if params['from_deg']==True:
            alpha = deg2rad(alpha)
            beta = deg2rad(beta)
            gamma = deg2rad(gamma)
        
        ca = math.cos(alpha)
        cb = math.cos(beta)
        cg = math.cos(gamma)
        sa = math.sin(alpha)
        sb = math.sin(beta)
        sg = math.sin(gamma)        

        Rx = np.asarray([[1, 0, 0], 
              [0, ca, sa], 
              [0, -sa, ca]
              ])

        Ry = np.asarray([[cb, 0, -sb], 
              [0, 1, 0],
              [sb, 0, cb]])

        Rz = np.asarray([[cg, sg, 0],
              [-sg, cg, 0],
              [0, 0, 1]])

        self.rotmat = np.eye(3)

        self.rotmat = np.matmul(Rx, self.rotmat)
        self.rotmat = np.matmul(Ry, self.rotmat)
        self.rotmat = np.matmul(Rz, self.rotmat)
        # self.rotmat = np.matmul(np.matmul(Rz, Ry), Rx)
   
    def _from_expmap(self, alpha, beta, gamma, params):
        if (alpha == 0 and beta == 0 and gamma == 0):
            self.rotmat = np.eye(3)
            return

        #TODO: Check exp map params

        theta = np.linalg.norm([alpha, beta, gamma])

        expmap = [alpha, beta, gamma] / theta

        x = expmap[0]
        y = expmap[1]
        z = expmap[2]

        s = np.sin(theta/2)
        c = np.cos(theta/2)

        self.rotmat = np.asarray([
            [2*(x**2-1)*s**2+1,  2*x*y*s**2-2*z*c*s,  2*x*z*s**2+2*y*c*s],
            [2*x*y*s**2+2*z*c*s,  2*(y**2-1)*s**2+1,  2*y*z*s**2-2*x*c*s],
            [2*x*z*s**2-2*y*c*s, 2*y*z*s**2+2*x*c*s , 2*(z**2-1)*s**2+1]
        ])
        


    def get_euler_axis(self):
        R = self.rotmat
        theta = np.arccos((self.rotmat.trace() - 1.0) * 0.5)
        axis = np.asarray([R[2,1] - R[1,2], R[0,2] - R[2,0], R[1,0] - R[0,1]])
        axis = 0.5 * axis/np.sin(theta)
        return theta, axis

    def get_euler_axis0(self):
        R = self.rotmat
        d, V = np.linalg.eig(R)        
        print(d.shape)
        print(V.shape)
        index = np.min(np.abs(np.diag(d) - 1))
        print(index)
        axis = V[:,index]
        
        cosTheta = (R.trace() - 1)/2
        _, index = max(abs(axis))

        if index == 0:            
            sinTheta = (R[2,1] - R[1,2])/(2*axis[index])
        elif index == 1:
            sinTheta = (R[0,2] - R[2,0])/(2*axis[index])                    
        elif index == 2:
            sinTheta = (R[1,0] - R[0,1])/(2*axis[index])
        
        theta = np.arctan2(sinTheta, cosTheta);

        return theta, axis

    def to_expmap0(self):
        theta, axis = self.get_euler_axis()
        rot_arr = theta * axis
        if np.isnan(rot_arr).any():
            rot_arr = [0, 0, 0]
        return rot_arr
    
    def to_expmap(self):
        R = self.rotmat
        return quat2expmap(rotmat2quat(R))

    def to_euler(self, use_deg=False):
        eulers = np.zeros((2, 3))

        if np.absolute(np.absolute(self.rotmat[2, 0]) - 1) < 1e-12:
            #GIMBAL LOCK!
            print('Gimbal')
            if np.absolute(self.rotmat[2, 0]) - 1 < 1e-12:
                eulers[:,0] = math.atan2(-self.rotmat[0,1], -self.rotmat[0,2])
                eulers[:,1] = -math.pi/2
            else:
                eulers[:,0] = math.atan2(self.rotmat[0,1], -self.rotmat[0,2])
                eulers[:,1] = math.pi/2
            
            return eulers

        theta = - math.asin(self.rotmat[2,0])
        theta2 = math.pi - theta

        # psi1, psi2
        eulers[0,0] = math.atan2(self.rotmat[2,1]/math.cos(theta), self.rotmat[2,2]/math.cos(theta))
        eulers[1,0] = math.atan2(self.rotmat[2,1]/math.cos(theta2), self.rotmat[2,2]/math.cos(theta2))

        # theta1, theta2
        eulers[0,1] = theta
        eulers[1,1] = theta2

        # phi1, phi2
        eulers[0,2] = math.atan2(self.rotmat[1,0]/math.cos(theta), self.rotmat[0,0]/math.cos(theta))
        eulers[1,2] = math.atan2(self.rotmat[1,0]/math.cos(theta2), self.rotmat[0,0]/math.cos(theta2))

        if use_deg:
            eulers = rad2deg(eulers)

        return eulers
    
    def to_quat(self):
        #TODO
        pass
    
    def __str__(self):
        return "Rotation Matrix: \n " + self.rotmat.__str__()
    



