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

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

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 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 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 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). 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). The rotation matrix representation of the constant rotation Get the inverse rotation, that is the rotation from the destination         reference frame to the source reference frame. 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. 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 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 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 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). 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). Get the inverse rotation, that is the rotation from the destination         reference frame to the source reference frame. 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 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 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. Apply the rotation to a position at a specific time 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.