# -*- coding: utf-8 -*-

# Copyright 2017, IBM.
#
# This source code is licensed under the Apache License, Version 2.0 found in
# the LICENSE.txt file in the root directory of this source tree.

"""
A module for using quaternions.
"""
import math
import numpy as np
import scipy.linalg as la


class Quaternion:
    """A class representing a Quaternion.
    """

    def __init__(self, data):
        self.data = np.asarray(data, dtype=float)

    def __call__(self, idx):
        return self.data[idx]

    def __repr__(self):
        return np.array_str(self.data)

    def __str__(self):
        return np.array_str(self.data)

    def __mul__(self, r):  # pylint: disable=C0103
        if isinstance(r, Quaternion):
            q = self  # pylint: disable=C0103
            out_data = np.zeros(4, dtype=float)
            out_data[0] = r(0)*q(0) - r(1)*q(1) - r(2)*q(2) - r(3)*q(3)
            out_data[1] = r(0)*q(1) + r(1)*q(0) - r(2)*q(3) + r(3)*q(2)
            out_data[2] = r(0)*q(2) + r(1)*q(3) + r(2)*q(0) - r(3)*q(1)
            out_data[3] = r(0)*q(3) - r(1)*q(2) + r(2)*q(1) + r(3)*q(0)
            return Quaternion(out_data)
        else:
            raise Exception('Multiplication by other not supported.')

    def norm(self):
        """ Norm of quaternion.
        """
        return la.norm(self.data)

    def normalize(self, inplace=False):
        """Normalizes a Quaternion to unit length
        so that it represents a valid roation.

        Args:
            inplace (bool): Do an inplace normalization.

        Returns:
            Quaternion: Normalized quaternion.
        """
        if inplace:
            nrm = self.norm()
            self.data /= nrm
            return None
        nrm = self.norm()
        data_copy = np.array(self.data, copy=True)
        data_copy /= nrm
        return Quaternion(data_copy)

    def to_matrix(self):
        """Converts a unit-length quaternion to a rotation matrix.

        Returns:
            ndarray: Rotation matrix.
        """
        w, x, y, z = self.normalize().data  # pylint: disable=C0103
        mat = np.array([
            [1-2*y**2-2*z**2, 2*x*y-2*z*w, 2*x*z+2*y*w],
            [2*x*y+2*z*w, 1-2*x**2-2*z**2, 2*y*z-2*x*w],
            [2*x*z-2*y*w, 2*y*z+2*x*w, 1-2*x**2-2*y**2]
        ], dtype=float)
        return mat

    def to_zyz(self):
        """Converts a unit-length quaternion to a sequence
        of ZYZ Euler angles.

        Returns:
            ndarray: Array of Euler angles.
        """
        mat = self.to_matrix()
        euler = np.zeros(3, dtype=float)
        if mat[2, 2] < 1:
            if mat[2, 2] > -1:
                euler[0] = math.atan2(mat[1, 2], mat[0, 2])
                euler[1] = math.acos(mat[2, 2])
                euler[2] = math.atan2(mat[2, 1], -mat[2, 0])
            else:
                euler[0] = -math.atan2(mat[1, 0], mat[1, 1])
                euler[1] = np.pi
        else:
            euler[0] = math.atan2(mat[1, 0], mat[1, 1])
        return euler


def quaternion_from_axis_rotation(angle, axis):
    """Return quaternion for rotation about given axis.

    Args:
        angle (float): Angle in radians.
        axis (str): Axis for rotation

    Returns:
        Quaternion: Quaternion for axis rotation.

    Raises:
        ValueError: Invalid input axis.
    """
    out = np.zeros(4, dtype=float)
    if axis == 'x':
        out[1] = 1
    elif axis == 'y':
        out[2] = 1
    elif axis == 'z':
        out[3] = 1
    else:
        raise ValueError('Invalid axis input.')
    out *= math.sin(angle/2.0)
    out[0] = math.cos(angle/2.0)
    return Quaternion(out)


def quaternion_from_euler(angles, order='yzy'):
    """Generate a quaternion from a set of Euler angles.

    Args:
        angles (array_like): Array of Euler angles.
        order (str): Order of Euler rotations.  'yzy' is default.

    Returns:
        Quaternion: Quaternion representation of Euler rotation.
    """
    angles = np.asarray(angles, dtype=float)
    quat = quaternion_from_axis_rotation(angles[0], order[0])\
        * (quaternion_from_axis_rotation(angles[1], order[1])
           * quaternion_from_axis_rotation(angles[2], order[2]))
    quat.normalize(inplace=True)
    return quat


def _rotm(angle, axis):
    """Generates a rotation matrix for a given angle and axis.

    Args:
        angle (float): Roation angle in radians.
        axis (str): Axis for rotation: 'x', 'y', 'z'

    Returns:
        ndarray: Rotation matrix.

    Raises:
        ValueError: Invalid input axis.
    """
    direction = np.zeros(3, dtype=float)
    if axis == 'x':
        direction[0] = 1
    elif axis == 'y':
        direction[1] = 1
    elif axis == 'z':
        direction[2] = 1
    else:
        raise ValueError('Invalid axis.')
    direction = np.asarray(direction, dtype=float)
    sin_angle = math.sin(angle)
    cos_angle = math.cos(angle)
    rot = np.diag([cos_angle, cos_angle, cos_angle])
    rot += np.outer(direction, direction) * (1.0 - cos_angle)
    direction *= sin_angle
    rot += np.array([
        [0, -direction[2], direction[1]],
        [direction[2], 0, -direction[0]],
        [-direction[1], direction[0], 0]
    ])
    return rot
