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

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 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)