# Copyright (c) 2015 Ultimaker B.V.
# Uranium is released under the terms of the AGPLv3 or higher.

import numpy
import numpy.linalg
import math
import copy

from UM.Math.Vector import Vector
from UM.Math.Float import Float
from UM.Math.Matrix import Matrix

##  Unit Quaternion class based on numpy arrays.
#
#   This class represents a Unit quaternion that can be used for rotations.
#
#   \note The operations that modify this quaternion will ensure the length
#         of the quaternion remains 1. This is done to make this class simpler
#         to use.
#
class Quaternion(object):
    EPS = numpy.finfo(float).eps * 4.0

    def __init__(self, x=0.0, y=0.0, z=0.0, w=1.0):
        # Components are stored as XYZW
        self._data = numpy.array([x, y, z, w], dtype=numpy.float32)

    def getData(self):
        return self._data

    @property
    def x(self):
        return self._data[0]

    @property
    def y(self):
        return self._data[1]

    @property
    def z(self):
        return self._data[2]

    @property
    def w(self):
        return self._data[3]

    ##  Set quaternion by providing rotation about an axis.
    #
    #   \param angle \type{float} Angle in radians
    #   \param axis \type{Vector} Axis of rotation
    def setByAngleAxis(self, angle, axis):
        a = axis.getNormalized().getData()
        halfAngle = angle / 2.0
        self._data[3] = math.cos(halfAngle)
        self._data[0:3] = a * math.sin(halfAngle)
        self.normalize()

    def __mul__(self, other):
        result = copy.deepcopy(self)
        result *= other
        return result

    def __imul__(self, other):
        if type(other) is Quaternion:
            v1 = Vector(other.x, other.y, other.z)
            v2 = Vector(self.x, self.y, self.z)

            w = other.w * self.w - v1.dot(v2)
            v = v2 * other.w + v1 * self.w + v2.cross(v1)

            self._data[0] = v.x
            self._data[1] = v.y
            self._data[2] = v.z
            self._data[3] = w
        elif type(other) is float or type(other) is int:
            self._data *= other
        else:
            raise NotImplementedError()

        return self

    def __add__(self, other):
        result = copy.deepcopy(self)
        result += other
        return result

    def __iadd__(self, other):
        if type(other) is Quaternion:
            self._data[0] += other._data[0]
            self._data[1] += other._data[1]
            self._data[2] += other._data[2]
            self._data[3] += other._data[3]
        else:
            raise NotImplementedError()

        return self

    def __truediv__(self, other):
        result = copy.deepcopy(self)
        result /= other
        return result

    def __itruediv__(self, other):
        if type(other) is float or type(other) is int:
            self._data /= other
        else:
            raise NotImplementedError()

        return self

    def __eq__(self, other):
        return Float.fuzzyCompare(self.x, other.x, 1e-6) and Float.fuzzyCompare(self.y, other.y, 1e-6) and Float.fuzzyCompare(self.z, other.z, 1e-6) and Float.fuzzyCompare(self.w, other.w, 1e-6)

    def __neg__(self):
        q = copy.deepcopy(self)
        q._data = -q._data
        return q

    def getInverse(self):
        result = copy.deepcopy(self)
        result.invert()
        return result

    def invert(self):
        self._data[0:3] = -self._data[0:3]
        return self

    def rotate(self, vector):
        vMult = 2.0 * (self.x * vector.x + self.y * vector.y + self.z * vector.z)
        crossMult = 2.0 * self.w
        pMult = crossMult * self.w - 1.0

        return Vector( pMult * vector.x + vMult * self.x + crossMult * (self.y * vector.z - self.z * vector.y),
                       pMult * vector.y + vMult * self.y + crossMult * (self.z * vector.x - self.x * vector.z),
                       pMult * vector.z + vMult * self.z + crossMult * (self.x * vector.y - self.y * vector.x) )

    def dot(self, other):
        return numpy.dot(self._data, other._data)

    def length(self):
        return numpy.linalg.norm(self._data)

    def normalize(self):
        self._data /= numpy.linalg.norm(self._data)

    ## Set quaternion by providing a homogenous (4x4) rotation matrix.
    # \param matrix 4x4 Matrix object
    # \param is_precise
    def setByMatrix(self, matrix, is_precise = False):
        trace = matrix.at(0, 0) + matrix.at(1, 1) + matrix.at(2, 2)
        if trace > 0.0:
            self._data[0] = matrix.at(2, 1) - matrix.at(1, 2)
            self._data[1] = matrix.at(0, 2) - matrix.at(2, 0)
            self._data[2] = matrix.at(1, 0) - matrix.at(0, 1)
            self._data[3] = trace + 1
        else:
            i = 0
            if matrix.at(1, 1) > matrix.at(0, 0):
                i = 1

            if matrix.at(2, 2) > matrix.at(i, i):
                i = 2

            # Yes, this is repeated code. Writing it out however makes the code way
            # more readable than any magical index shifting.
            if i == 0:
                self._data[0] = matrix.at(0, 0) - matrix.at(1, 1) - matrix.at(2, 2) + 1.0
                self._data[1] = matrix.at(0, 1) + matrix.at(1, 0)
                self._data[2] = matrix.at(0, 2) + matrix.at(2, 0)
                self._data[3] = matrix.at(2, 1) - matrix.at(1, 2)
            elif i == 1:
                self._data[0] = matrix.at(0, 1) + matrix.at(1, 0)
                self._data[1] = matrix.at(1, 1) - matrix.at(0, 0) - matrix.at(2, 2) + 1.0
                self._data[2] = matrix.at(1, 2) + matrix.at(2, 1)
                self._data[3] = matrix.at(0, 2) - matrix.at(2, 0)
            else:
                self._data[0] = matrix.at(0, 2) + matrix.at(2, 0)
                self._data[1] = matrix.at(2, 1) + matrix.at(1, 2)
                self._data[1] = matrix.at(2, 2) - matrix.at(0, 0) - matrix.at(1, 1) + 1.0
                self._data[3] = matrix.at(1, 0) - matrix.at(0, 1)

        self.normalize()

    def toMatrix(self):
        m = numpy.zeros((4, 4), dtype=numpy.float32)

        s = 2.0 / (self.x ** 2 + self.y ** 2 + self.z ** 2 + self.w ** 2)

        xs = s * self.x
        ys = s * self.y
        zs = s * self.z

        wx = self.w * xs
        wy = self.w * ys
        wz = self.w * zs

        xx = self.x * xs
        xy = self.x * ys
        xz = self.x * zs

        yy = self.y * ys
        yz = self.y * zs
        zz = self.z * zs

        m[0,0] = 1.0 - (yy + zz)
        m[0,1] = xy - wz
        m[0,2] = xz + wy

        m[1,0] = xy + wz
        m[1,1] = 1.0 - (xx + zz)
        m[1,2] = yz - wx

        m[2,0] = xz - wy
        m[2,1] = yz + wx
        m[2,2] = 1.0 - (xx + yy)

        m[3,3] = 1.0

        return Matrix(m)

    @staticmethod
    def slerp(start, end, amount):
        if Float.fuzzyCompare(amount, 0.0):
            return start
        elif Float.fuzzyCompare(amount, 1.0):
            return end

        rho = math.acos(start.dot(end))
        return (start * math.sin((1 - amount) * rho) + end * math.sin(amount * rho)) / math.sin(rho)

    ##  Returns a quaternion representing the rotation from vector 1 to vector 2.
    #
    #   \param v1 \type{Vector} The vector to rotate from.
    #   \param v2 \type{Vector} The vector to rotate to.
    @staticmethod
    def rotationTo(v1, v2):
        d = v1.dot(v2)

        if d >= 1.0:
            return Quaternion() # Vectors are equal, no rotation needed.

        q = None
        if Float.fuzzyCompare(d, -1.0, 1e-6):
            axis = Vector.Unit_X.cross(v1)

            if Float.fuzzyCompare(axis.length(), 0.0):
                axis = Vector.Unit_Y.cross(v1)

            axis.normalize()
            q = Quaternion()
            q.setByAngleAxis(math.pi, axis)
        else:
            s = math.sqrt((1.0 + d) * 2.0)
            invs = 1.0 / s

            c = v1.cross(v2)

            q = Quaternion(
                c.x * invs,
                c.y * invs,
                c.z * invs,
                s * 0.5
            )
            q.normalize()

        return q

    @staticmethod
    def fromMatrix(matrix):
        q = Quaternion()
        q.setByMatrix(matrix)
        return q

    @staticmethod
    def fromAngleAxis(angle, axis):
        q = Quaternion()
        q.setByAngleAxis(angle, axis)
        return q

    def __repr__(self):
        return "Quaternion(x={0}, y={1}, z={2}, w={3})".format(self.x, self.y, self.z, self.w)