from scipy.interpolate import interp1d
from scipy.spatial.transform import Rotation

import numpy as np

class ConstantRotation:
    """
    A constant rotation between two 3D reference frames.

    Attributes
    __________
    source : int
             The NAIF ID code for the source frame
    dest : int
           The NAIF ID code for the destination frame
    """

    def from_matrix(mat, source, dest):
        """
        Create a constant rotation from a directed cosine matrix

        Parameters
        ----------
        mat : 2darray
              The rotation matrix
        source : int
                 The NAIF ID code for the source frame
        dest : int
               The NAIF ID code for the destination frame

        See Also
        --------
        scipy.spatial.transform.Rotation.from_dcm
        """
        rot = Rotation.from_dcm(mat)
        return ConstantRotation(rot.as_quat(), source, dest)

    def __init__(self, quat, source, dest):
        """
        Construct a constant rotation

        Parameters
        ----------
        quat : array
               The quaternion representation of the rotation as a numpy array.
               The quaternion must be in scalar last format (x, y, z, w).
        source : int
                 The NAIF ID code for the source frame
        dest : int
               The NAIF ID code for the destination frame
        """
        self.source = source
        self.dest = dest
        self.quat = np.asarray(quat)

    def __repr__(self):
        return f'ConstantRotation Source: {self.source}, Destination: {self.dest}, Quat: {self.quat}'

    @property
    def quat(self):
        """
        The quaternion that rotates from the source reference frame to
        the destination reference frame. The quaternion is in scalar last
        format (x, y, z, w).
        """
        return self._rot.as_quat()

    @quat.setter
    def quat(self, new_quat):
        """
        Change the rotation to a different quaternion

        Parameters
        ----------
        new_quat : array
                   The new quaternion as an array.
                   The quaternion must be in scalar last format (x, y, z, w).
        """
        self._rot = Rotation.from_quat(np.asarray(new_quat))

    def rotation_matrix(self):
        """
        The rotation matrix representation of the constant rotation
        """
        return self._rot.as_dcm()

    def inverse(self):
        """
        Get the inverse rotation, that is the rotation from the destination
        reference frame to the source reference frame.
        """
        return ConstantRotation(self._rot.inv().as_quat(), self.dest, self.source)

    def __mul__(self, other):
        """
        Compose this rotation with another rotation.

        The destination frame of the right rotation (other) and the source
        frame of the left rotation (self) must be the same. I.E. if A and B are
        rotations, then for A*B to be valid, A.source must equal B.dest.

        Parameters
        ----------
        other : Rotation
                Another rotation object, it can be constant or time dependent.
        """
        if self.source != other.dest:
            raise ValueError("Destination frame of first rotation {} is not the same as source frame of second rotation {}.".format(other.dest, self.source))
        if isinstance(other, ConstantRotation):
            new_rot = self._rot * other._rot
            return ConstantRotation(new_rot.as_quat(), other.source, self.dest)
        elif isinstance(other, TimeDependentRotation):
            return TimeDependentRotation((self._rot * other._rots).as_quat(), other.times, other.source, self.dest)
        else:
            raise TypeError("Rotations can only be composed with other rotations.")

class TimeDependentRotation:
    """
    A time dependent rotation between two 3D reference frames.

    Attributes
    __________
    source : int
             The NAIF ID code for the source frame
    dest : int
           The NAIF ID code for the destination frame
    """

    def from_euler(sequence, euler, times, source, dest, degrees=False):
        """
        Create a time dependent rotation from a set of Euler angles.

        Parameters
        ----------
        sequence : string
                   The axis sequence that the Euler angles are applied in. I.E. 'XYZ'
                   or 'ZXZ'.
        euler : 2darray
                2D numpy array of the euler angle rotations in radians.
        times : array
                The time for each rotation in euler. This array must be sorted
                in ascending order.
        source : int
                 The NAIF ID code for the source frame
        dest : int
               The NAIF ID code for the destination frame
        degrees : bool
                  If the angles are in degrees. If false, then degrees are
                  assumed to be in radians. Defaults to False.

        See Also
        --------
        scipy.spatial.transform.Rotation.from_euler
        """
        rot = Rotation.from_euler(sequence, np.asarray(euler), degrees=degrees)
        return TimeDependentRotation(rot.as_quat(), times, source, dest)

    def __init__(self, quats, times, source, dest):
        """
        Construct a time dependent rotation

        Parameters
        ----------
        quats : 2darray
               The quaternion representations of the rotation as a 2d numpy array.
               Each inner array represents the rotation at the time at the same index
               in the times argument. The quaternions must be in scalar last format
               (x, y, z, w).
        times : array
                The time for each rotation in quats. This array must be sorted
                in ascending order.
        source : int
                 The NAIF ID code for the source frame
        dest : int
               The NAIF ID code for the destination frame
        """
        self.source = source
        self.dest = dest
        self.quats = np.asarray(quats)
        self.times = np.asarray(times)

    def __repr__(self):
        return f'Time Dependent Rotation Source: {self.source}, Destination: {self.dest}, Quat: {self.quats}'

    @property
    def quats(self):
        """
        The quaternions that rotates from the source reference frame to
        the destination reference frame. The quaternions are in scalar
        last format (x, y, z, w).
        """
        return self._rots.as_quat()

    @quats.setter
    def quats(self, new_quats):
        """
        Change the rotations to interpolate over

        Parameters
        ----------
        new_quats : 2darray
                    The new quaternions as a 2d array. The quaternions must be
                    in scalar last format (x, y, z, w).
        """
        self._rots = Rotation.from_quat(np.asarray(new_quats))

    def inverse(self):
        """
        Get the inverse rotation, that is the rotation from the destination
        reference frame to the source reference frame.
        """
        return TimeDependentRotation(self._rots.inv().as_quat(), self.times, self.dest, self.source)

    def _slerp(self, times):
        """
        Using SLERP interpolate the rotation and angular velocity at specific times

        This uses the same code as scipy SLERP, except it extrapolates
        assuming constant angular velocity before and after the first
        and last intervals.

        If the this rotation is only defined at one time, then the rotation is
        assumed to be constant.

        Parameters
        ----------
        times : 1darray or float
                The new times to interpolate at.

        Returns
        -------
         : Rotation
           The new rotations at the input times
         : 2darray
           The angular velocity vectors
        """
        vec_times = np.asarray(times)
        if vec_times.ndim < 1:
            vec_times = np.asarray([times])
        elif vec_times.ndim > 1:
            raise ValueError('Input times must be either a float or a 1d iterable of floats')
        if len(self.times) < 2:
            return Rotation.from_quat(np.repeat(self.quats, len(vec_times), 0)), np.zeros((len(vec_times), 3))
        else:
            idx = np.searchsorted(self.times, vec_times) - 1
            idx[idx >= len(self.times) - 1] = len(self.times) - 2
            idx[idx < 0] = 0
            steps = self.times[idx+1] - self.times[idx]
            rotvecs = (self._rots[idx + 1] * self._rots[idx].inv()).as_rotvec()
            alpha = (vec_times - self.times[idx]) / steps
            interp_rots = Rotation.from_rotvec(rotvecs * alpha[:, None]) * self._rots[idx]
            interp_av = rotvecs / steps[:, None]
            return interp_rots, interp_av

    def reinterpolate(self, times):
        """
        Reinterpolate the rotation at a given set of times.

        Parameters
        ----------
        times : 1darray or float
                The new times to interpolate at.

        Returns
        -------
         : TimeDependentRotation
           The new rotation that the input times
        """
        new_rots, _ = self._slerp(times)
        return TimeDependentRotation(new_rots.as_quat(), times, self.source, self.dest)

    def __mul__(self, other):
        """
        Compose this rotation with another rotation.

        The destination frame of the right rotation (other) and the source
        frame of the left rotation (self) must be the same. I.E. if A and B are
        rotations, then for A*B to be valid, A.source must equal B.dest.

        If the other rotation is a time dependent rotation, then the time range
        for the resultant rotation will be the time covered by both rotations.
        I.E. if A covers 0 to 2 and B covers 1 to 4, then A*B will cover 1 to 2.

        Parameters
        ----------
        other : Rotation
                Another rotation object, it can be constant or time dependent.
        """
        if self.source != other.dest:
            raise ValueError("Destination frame of first rotation {} is not the same as source frame of second rotation {}.".format(other.dest, self.source))
        if isinstance(other, ConstantRotation):
            return TimeDependentRotation((self._rots * other._rot).as_quat(), self.times, other.source, self.dest)
        elif isinstance(other, TimeDependentRotation):
            merged_times = np.union1d(np.asarray(self.times), np.asarray(other.times))
            new_quats = (self.reinterpolate(merged_times)._rots * other.reinterpolate(merged_times)._rots).as_quat()
            return TimeDependentRotation(new_quats, merged_times, other.source, self.dest)
        else:
            raise TypeError("Rotations can only be composed with other rotations.")

    def apply_at(self, vec, et):
        """
        Apply the rotation to a position at a specific time
        """
        return self.reinterpolate(et)._rots.apply(vec)

    def rotate_velocity_at(self, pos, vel, et):
        """
        Apply the rotation to a velocity at a specific time

        See:
        https://naif.jpl.nasa.gov/pub/naif/toolkit_docs/C/req/rotation.html#State%20transformations
        For an explanation of why a separate method is required to rotate velocities.
        """
        vec_pos = np.asarray(pos)
        vec_vel = np.asarray(vel)
        if vec_pos.ndim < 1:
            vec_pos = np.asarray([pos])
        if vec_vel.ndim < 1:
            vec_vel = np.asarray([vel])
        if vec_pos.shape != vec_vel.shape:
            raise ValueError('Input velocities and positions must have the same shape')

        rots, avs = self._slerp(et)
        rotated_vel = np.zeros(vec_vel.shape)
        for indx in range(vec_pos.shape[0]):
            skew = np.array([[0, -avs[indx, 2], avs[indx, 1]],
                             [avs[indx, 2], 0, -avs[indx, 0]],
                             [-avs[indx, 1], avs[indx, 0], 0]])
            rot_deriv = np.dot(skew, rots[indx].as_dcm())
            rotated_vel[indx] = rots[indx].apply(vec_vel[indx])
            rotated_vel[indx] += np.dot(rot_deriv, vec_pos[indx])

        return rotated_vel
