#!/usr/bin/env python3
# -*- coding: utf-8 -*-

from .Message import Message

import math
import numpy as np

class Quaternion(Message):
    """
    Basic class for representing and operation with quaternions that represent orientation in the free space
    Data entries:
    x,y,z,w component of the quaternion
    """
    __attributes__ = ['x', 'y','z','w']
    __attribute_types = ['float32','float32','float32','float32']
    
    def __init__(self, *args, **kwds):
        super(Quaternion, self).__init__(*args, **kwds)
        #message fields cannot be None 
        if self.x == None or self.y == None or self.z == None or self.w == None:
            self.x = 0
            self.y = 0
            self.z = 0
            self.w = 1
    
    def to_R(self):
        """Method to convert quaternion into rotation matrix
        Returns:
            rotation matrix
        """
        q = [self.w, self.x, self.y, self.z]
        R = [[q[0]**2+q[1]**2-q[2]**2-q[3]**2,     2*(q[1]*q[2]-q[0]*q[3]),      2*(q[1]*q[3]+q[0]*q[2])],
             [2*(q[1]*q[2]+q[0]*q[3]),     q[0]**2-q[1]**2+q[2]**2-q[3]**2,      2*(q[2]*q[3]-q[0]*q[1])],
             [2*(q[1]*q[3]-q[0]*q[2]),         2*(q[2]*q[3]+q[0]*q[1]),   q[0]**2-q[1]**2-q[2]**2+q[3]**2]]
        return np.asarray(R)

    def from_R(self, R):
        """Method to convert rotation matrix into quaternion
        Args:
            R: rotation matrix
        """
        q = np.array([R.trace() + 1, R[2,1]-R[1,2],R[0,2]-R[2,0],R[1,0]-R[0,1]])
        q = 1.0/(2.0*math.sqrt(R.trace() + 1.0000001))*q
        q = q/np.linalg.norm(q)
        self.x = q[1]
        self.y = q[2]
        self.z = q[3]
        self.w = q[0]

    def to_Euler(self):
        """Method to convert quaternion into euler angles  
        Returns:
            yaw, pitch, roll Euler angles
        """
        t0 = +2.0 * (self.w * self.x + self.y * self.z)
        t1 = +1.0 - 2.0 * (self.x * self.x + self.y * self.y)
        roll = math.atan2(t0, t1)
        t2 = +2.0 * (self.w * self.y - self.z * self.x)
        t2 = +1.0 if t2 > +1.0 else t2
        t2 = -1.0 if t2 < -1.0 else t2
        pitch = math.asin(t2)
        t3 = +2.0 * (self.w * self.z + self.x * self.y)
        t4 = +1.0 - 2.0 * (self.y * self.y + self.z * self.z)
        yaw = math.atan2(t3, t4)
        return yaw, pitch, roll
    
    def from_Euler(self, yaw, pitch, roll):
        """Method to convert euler angles into quaternion
        Args:
            yaw, pitch, roll: Euler angles
        """
        q = np.zeros(4)
        q[0] = np.cos(roll/2) * np.cos(pitch/2) * np.cos(yaw/2) + np.sin(roll/2) * np.sin(pitch/2) * np.sin(yaw/2)
        q[1] = np.sin(roll/2) * np.cos(pitch/2) * np.cos(yaw/2) - np.cos(roll/2) * np.sin(pitch/2) * np.sin(yaw/2)
        q[2] = np.cos(roll/2) * np.sin(pitch/2) * np.cos(yaw/2) + np.sin(roll/2) * np.cos(pitch/2) * np.sin(yaw/2)
        q[3] = np.cos(roll/2) * np.cos(pitch/2) * np.sin(yaw/2) - np.sin(roll/2) * np.sin(pitch/2) * np.cos(yaw/2)
        q = q/np.linalg.norm(q)
        self.x = q[1]
        self.y = q[2]
        self.z = q[3]
        self.w = q[0]


#Examples of usage
if __name__=="__main__":
    q = Quaternion()
    print("quaternion\n", q.x, q.y, q.z, q.w)
    print("rotation matrix\n", q.to_R())

    #given by euler angles
    yaw = math.pi 
    pitch = 0 
    roll = 0 
    q = Quaternion()
    q.from_Euler(yaw, pitch, roll)
    print("quaternion\n", q.x, q.y, q.z, q.w)
    #show as rotation matrix
    print("rotation matrix\n", q.to_R())
