import numpy as np

class QuinticTrajectory:
    """五次多项式轨迹规划"""
    
    def __init__(self, q0, qf, v0=0, vf=0, a0=0, af=0, T=1.0):
        self.q0 = q0
        self.qf = qf
        self.v0 = v0
        self.vf = vf
        self.a0 = a0
        self.af = af
        self.T = T
        self._compute_coeff()
    
    def _compute_coeff(self):
        T = self.T
        q0, qf = self.q0, self.qf
        v0, vf = self.v0, self.vf
        a0, af = self.a0, self.af
        
        self.a0 = q0
        self.a1 = v0
        self.a2 = a0 / 2.0
        
        A = np.array([
            [T**3, T**4, T**5],
            [3*T**2, 4*T**3, 5*T**4],
            [6*T, 12*T**2, 20*T**3]
        ])
        b = np.array([
            qf - q0 - v0*T - 0.5*a0*T**2,
            vf - v0 - a0*T,
            af - a0
        ])
        coeff = np.linalg.solve(A, b)
        self.a3, self.a4, self.a5 = coeff
    
    def get_position(self, t):
        t = np.clip(t, 0, self.T)
        return (self.a0 + self.a1*t + self.a2*t**2 + 
                self.a3*t**3 + self.a4*t**4 + self.a5*t**5)
    
    def get_velocity(self, t):
        t = np.clip(t, 0, self.T)
        return (self.a1 + 2*self.a2*t + 3*self.a3*t**2 + 
                4*self.a4*t**3 + 5*self.a5*t**4)
    
    def get_acceleration(self, t):
        t = np.clip(t, 0, self.T)
        return (2*self.a2 + 6*self.a3*t + 
                12*self.a4*t**2 + 20*self.a5*t**3)