'''
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

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)

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

        if params['from_deg']==True:
            alpha = deg2rad(alpha)
            beta = deg2rad(beta)
            gamma = deg2rad(gamma)
        
        Rx = np.asarray([[1, 0, 0], 
              [0, math.cos(alpha), -math.sin(alpha)], 
              [0, math.sin(alpha), math.cos(alpha)]
              ])

        Ry = np.asarray([[math.cos(beta), 0, math.sin(beta)], 
              [0, 1, 0],
              [-math.sin(beta), 0, math.cos(beta)]])

        Rz = np.asarray([[math.cos(gamma), -math.sin(gamma), 0],
              [math.sin(gamma), math.cos(gamma), 0],
              [0, 0, 1]])

        self.rotmat = np.matmul(np.matmul(Rz, Ry), Rx).T
    
    def get_euler_axis(self):
        R = self.rotmat
        theta = math.acos((self.rotmat.trace() - 1) / 2)
        axis = np.asarray([R[2,1] - R[1,2], R[0,2] - R[2,0], R[1,0] - R[0,1]])
        axis = axis/(2*math.sin(theta))
        return theta, axis

    def to_expmap(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_euler(self):
        #TODO
        pass
    
    def to_quat(self):
        #TODO
        pass