import numpy as np
from scipy.interpolate import CubicSpline
from threading import Lock

def change_in_pos(v0, a0, jerk, t1, t2):
    """
    jerk = 0 if constant acceleration.
    """
    return v0 * (t2 - t1) + a0 * (t2 - t1)**2 / 2 + jerk * (t2 - t1)**3 / 6

def change_in_vel(a0, jerk, t1, t2):
    return a0 * (t2 - t1) + jerk * (t2 - t1)**2 / 2

def create_cubic_spline(min_value, max_value, num_steps):
    x_points = [0, 1]
    
    y_points = [min_value, max_value]
    cs = CubicSpline(x_points, y_points, bc_type='clamped')
    x_steps = np.linspace(0, 1, num_steps)
    y_steps = cs(x_steps)
    return y_steps

def clamp(x, lo, hi):
    return lo if x < lo else hi if x > hi else x

def create_cubic_hermite(p0, v0, p1, v1, num_steps):
    """
    Cubic Hermite interpolation from p0->p1 with endpoint slopes v0, v1
    over num_steps samples (including endpoints).
    v0 and v1 are in "value per step" units.
    """
    if num_steps <= 1:
        return [p1]

    traj = []
    n = num_steps - 1  # steps between endpoints
    for i in range(num_steps):
        t = i / n  # 0..1
        t2 = t * t
        t3 = t2 * t

        h00 = 2 * t3 - 3 * t2 + 1
        h10 = t3 - 2 * t2 + t
        h01 = -2 * t3 + 3 * t2
        h11 = t3 - t2

        # Scale slopes by total duration (n) to match Hermite formulation.
        # Here v0, v1 are per-step, so multiply by n.
        pt = (h00 * p0) + (h10 * (v0 * n)) + (h01 * p1) + (h11 * (v1 * n))
        traj.append(pt)
    return traj

class Trajectory:
    """
    A trajectory is a list of values that are used to control a motor.
    It is used to smoothly transition the motor from one position to another.

    Trajectories is based around calibrated_angle because we want to make it readable around the URDF.
    """

    def __init__(self, motor):
        self.index = 0
        self.motor = motor
        self.trajectory = []
        self.trajectory_lock = Lock()

        # Tunables
        self.max_slope = None       # e.g. set to a float to cap velocity changes per step
        self.end_slope = 0.0        # slope at the end (0 gives a gentle settle)

    def _current_value(self):
        # Prefer in-flight trajectory point if we have one, otherwise motor target.
        if self.trajectory and self.index < len(self.trajectory):
            return self.trajectory[self.index]
        return self.motor.convert_angle_to_calibrated_angle(self.motor.target_position)

    def _estimate_current_slope(self):
        """
        Estimate per-step slope from the existing trajectory near the current index.
        This gives continuity of velocity when retargeting mid-move.
        """
        if not self.trajectory:
            return 0.0

        i = self.index
        # Use a central/forward difference when possible
        if 1 <= i < len(self.trajectory) - 1:
            v = 0.5 * (self.trajectory[i + 1] - self.trajectory[i - 1])
        elif 0 <= i < len(self.trajectory) - 1:
            v = self.trajectory[i + 1] - self.trajectory[i]
        elif len(self.trajectory) >= 2:
            v = self.trajectory[-1] - self.trajectory[-2]
        else:
            v = 0.0

        if self.max_slope is not None:
            v = clamp(v, -self.max_slope, self.max_slope)
        return v

    def update_trajectory(self, end_value, num_steps):
        """Update the trajectory from the current position to go to the end_value over num_steps.

        Args:
            end_value: The target calibrated_angle to reach.
            num_steps: The number of steps to take to reach the end_value.
        """
        with self.trajectory_lock:
            p0 = self._current_value()
            v0 = self._estimate_current_slope()

            # Optional: also cap v1 (end slope)
            v1 = self.end_slope
            if self.max_slope is not None:
                v1 = clamp(v1, -self.max_slope, self.max_slope)

            self.index = 0
            self.trajectory = create_cubic_hermite(p0, v0, end_value, v1, num_steps)

    def step(self):
        with self.trajectory_lock:
            if len(self.trajectory) == 0:
                return self.motor.calibrated_angle[0]
            elif len(self.trajectory) >= self.index + 1:
                to_return = self.trajectory[self.index]
                self.index += 1
                return to_return
            else:
                return self.trajectory[-1]

    def reset(self):
        with self.trajectory_lock:
            self.index = 0
            self.trajectory = []

    def is_at_end(self):
        return self.index >= len(self.trajectory)

# class Trajectory:
#     """
#     A trajectory is a list of values that are used to control a motor.
#     It is used to smoothly transition the motor from one position to another.

#     Trajectories is based around calibrated_angle because we want to make it readable around the URDF.
#     """

#     def __init__(self, motor):
#         self.index = 0
#         self.motor = motor
#         self.trajectory = []
#         self.trajectory_lock = Lock()

#     def update_trajectory(self, end_value, num_steps):
#         """Update the trajectory from the current position to go to the end_value over num_steps.

#         Args:
#             end_value: The target calibrated_angle to reach.
#             num_steps: The number of steps to take to reach the end_value.
#         """
#         with self.trajectory_lock:
#             self.index = 0
#             self.trajectory = create_cubic_spline(self.motor.convert_angle_to_calibrated_angle(self.motor.target_position), end_value, num_steps)

#     def add_trajectory(self, end_value, num_steps):
#         """Add a trajectory to the current trajectory.

#         Args:
#             end_value: The target calibrated_angle to reach.
#             num_steps: The number of steps to take to reach the end_value.
#         """
#         with self.trajectory_lock:
#             self.trajectory = np.concatenate([self.trajectory, create_cubic_spline(self.trajectory[-1], end_value, num_steps)])

#     def step(self):
#         with self.trajectory_lock:
#             if len(self.trajectory) == 0:
#                 return self.motor.calibrated_angle[0]
#             elif len(self.trajectory) >= self.index + 1:
#                 to_return = self.trajectory[self.index]
#                 self.index += 1
#                 return to_return
#             else:
#                 return self.trajectory[-1]

#     def reset(self):
#         with self.trajectory_lock:
#             self.index = 0
#             self.trajectory = []

#     def is_at_end(self):
#         return self.index >= len(self.trajectory)