from __future__ import absolute_import, print_function

import numpy as np

class Quaternion(object):
    def __init__(self,w=1.,x=0.,y=0.,z=0.):
        
        self.data = np.array([w,x,y,z])

    @classmethod
    def copy(cls,rhs):
        return Quaternion(*rhs.data)


    def __getitem__(self,i):
        return self.data[i]

    def __setitem__(self,i,val):
        self.data[i] = val



    def conj(self):
        return Quaternion(self[0],-self[1],-self[2],-self[2])

    def norm(self):
        return np.linalg.norm(self.data)

    def __add__(self,q):
        return Quaternion(*(self.data+q.data))

    def __sub__(self,q):
        return Quaternion(*(self.data-q.data))


    def __mul__(self,q):
        if isinstance(q,Quaternion):
            a1,b1,c1,d1 = self.data
            a2,b2,c2,d2 = q.data
            return Quaternion(a1*a2 - b1*b2 - c1*c2 - d1*d2,
                          a1*b2 + b1*a2 + c1*d2 - d1*c2,
                          a1*c2 - b1*d2 + c1*a2 + d1*b2,
                          a1*d2 + b1*c2 - c1*b2 + d1*a2)
        else:
            return Quaternion(*(q*self.data))

    def __repr__(self):
        return "Quaternion(%s,%s,%s,%s)"%tuple(self.data)

    def dot(self,q):
        return np.inner(self.data,q.data)

    def normalize(self):
        return Quaternion(*(self.data*1./self.norm()))

    def toRotation4(self):
        a,b,c,d = self.data
        return np.array([
    [a**2+b**2-c**2-d**2, 2*(b*c-a*d), 2*(b*d+a*c),0],
    [2*(b*c+a*d), a**2-b**2+c**2-d**2, 2*(c*d-a*b),0],
    [2*(b*d-a*c), 2*(c*d+a*b),  a**2-b**2-c**2+d**2,0],
    [0,0,0,1]
    ])

    def toRotation3(self):
        a,b,c,d = self.data
        return np.array([
    [a**2+b**2-c**2-d**2, 2*(b*c-a*d), 2*(b*d+a*c)],
    [2*(b*c+a*d), a**2-b**2+c**2-d**2, 2*(c*d-a*b)],
    [2*(b*d-a*c), 2*(c*d+a*b),  a**2-b**2-c**2+d**2],
    ])



def quaternion_slerp(q1,q2,t):
    q1,q2 = q1.normalize(), q2.normalize()

    prod = q1.dot(q2)
    if abs(prod)>.9998:
        return q1+(q2-q1)*t
    #picks the shorter great circle
    if prod<0:
        q2 = q2*(-1.)
        prod *= -1.

    w = np.arccos(prod)
    return (q1*(np.sin((1.-t)*w)/np.sin(w)))+q2*(np.sin(t*w)/np.sin(w))

if __name__ == '__main__':
    q1 = Quaternion(1,0,0,0)
    q2 = Quaternion(0,1,0,0)

    for t in np.linspace(0,1,10):
        q = quaternion_slerp(q1,q2,t)
        print(t, q, q.norm())
