import numpy as np
from scipy.interpolate import CubicSpline

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:
    def __init__(self, motor):
        self.index = 0
        self.motor = motor
        self.trajectory = []

    def update_trajectory(self, end_value, num_steps):
        self.index = 0
        self.trajectory = create_cubic_spline(self.motor.target_position, end_value, num_steps)

    def step(self):
        if len(self.trajectory) == 0:
            return self.motor.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):
        self.index = 0
        self.trajectory = []

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