import numpy as np
from scipy.spatial.transform import Rotation as R, Slerp

class TrajectoryInterpolator:
    def __init__(self, waypoints):
        self.waypoints = sorted(waypoints, key=lambda x: x["time"])
        self.times = [wp["time"] for wp in self.waypoints]
        self._compute_position_coeffs()
        self._prepare_slerp()

    def _compute_position_coeffs(self):
        self.coeffs_pos = []
        for i in range(len(self.waypoints) - 1):
            wp0 = self.waypoints[i]
            wp1 = self.waypoints[i + 1]
            T = wp1["time"] - wp0["time"]

            p0 = np.array(wp0["position"])
            v0 = np.array(wp0.get("velocity", np.zeros(3)))
            a0 = np.array(wp0.get("acceleration", np.zeros(3)))

            p1 = np.array(wp1["position"])
            v1 = np.array(wp1.get("velocity", np.zeros(3)))
            a1 = np.array(wp1.get("acceleration", np.zeros(3)))

            # 系数矩阵 A（以 tau = 0~T）
            A = np.array([
                [1, 0, 0,    0,     0,      0],
                [0, 1, 0,    0,     0,      0],
                [0, 0, 2,    0,     0,      0],
                [1, T, T**2, T**3,  T**4,   T**5],
                [0, 1, 2*T,  3*T**2,4*T**3, 5*T**4],
                [0, 0, 2,    6*T,  12*T**2, 20*T**3]
            ])

            b = np.vstack([p0, v0, a0, p1, v1, a1])  # shape (6,3)
            coeffs = np.linalg.solve(A, b)  # shape (6,3), 每列是一个维度的系数
            self.coeffs_pos.append((coeffs, T))

    def _prepare_slerp(self):
        self.rot_slerps = []
        for i in range(len(self.waypoints) - 1):
            t0 = self.waypoints[i]["time"]
            t1 = self.waypoints[i + 1]["time"]
            rot0 = R.from_quat(self.waypoints[i]["orientation"])
            rot1 = R.from_quat(self.waypoints[i + 1]["orientation"])
            self.rot_slerps.append(Slerp([0, 1], R.concatenate([rot0, rot1])))

    def interpolate(self, t):
        if t <= self.times[0]:
            return self._format_output(self.waypoints[0])
        elif t >= self.times[-1]:
            return self._format_output(self.waypoints[-1])

        i = np.searchsorted(self.times, t, side='right') - 1
        tau = t - self.times[i]
        coeffs, T = self.coeffs_pos[i]
        alpha = tau / T

        # 每个维度分开处理
        pos = np.array([np.polyval(coeffs[:, dim][::-1], tau) for dim in range(3)])
        vel = np.array([np.polyval(np.polyder(coeffs[:, dim][::-1], 1), tau) for dim in range(3)])
        acc = np.array([np.polyval(np.polyder(coeffs[:, dim][::-1], 2), tau) for dim in range(3)])

        # 姿态 slerp
        ori = self.rot_slerps[i]([alpha])[0].as_quat()

        ang_vel, ang_acc = self._estimate_angular_derivatives(t)

        x_r = np.concatenate([pos, ori])
        v_r = np.concatenate([vel, ang_vel])
        a_r = np.concatenate([acc, ang_acc])
        return x_r, v_r, a_r

    def _format_output(self, wp):
        pos = np.array(wp["position"])
        ori = np.array(wp["orientation"])
        vel = np.array(wp.get("velocity", np.zeros(3)))
        acc = np.array(wp.get("acceleration", np.zeros(3)))
        x_r = np.concatenate([pos, ori])
        v_r = np.concatenate([vel, np.zeros(3)])
        a_r = np.concatenate([acc, np.zeros(3)])
        return x_r, v_r, a_r
    def _quat_diff_to_angular_velocity(self, q1, q0, dt):
        """计算从 q0 到 q1 的角速度，单位 rad/s"""
        dq = R.from_quat(q1) * R.from_quat(q0).inv()
        angle = dq.magnitude()
        axis = dq.as_rotvec()
        if dt == 0 or angle == 0:
            return np.zeros(3)
        return axis / dt  # rotvec 本身是 axis * angle

    def _estimate_angular_derivatives(self, t, delta=0.03):
        """数值估计角速度和角加速度，避免递归调用"""
        if t <= self.times[0] or t >= self.times[-1]:
            return np.zeros(3), np.zeros(3)

        i = np.searchsorted(self.times, t, side='right') - 1
        t0, t1 = self.times[i], self.times[i + 1]
        T = t1 - t0
        alpha = (t - t0) / T

        alpha_prev = max(0.0, (t - delta - t0) / T)
        alpha_next = min(1.0, (t + delta - t0) / T)

        q_prev = self.rot_slerps[i]([alpha_prev])[0].as_quat()
        q_now = self.rot_slerps[i]([alpha])[0].as_quat()
        q_next = self.rot_slerps[i]([alpha_next])[0].as_quat()

        w1 = self._quat_diff_to_angular_velocity(q_now, q_prev, delta)
        w2 = self._quat_diff_to_angular_velocity(q_next, q_now, delta)
        w = w1
        a = (w2 - w1) / (2 * delta)
        return w, a


    
if __name__ == "__main__":
    import matplotlib.pyplot as plt

    waypoints = [
        {"time": 0.0, "position": [0, 0, 0], "velocity": [0, 0, 0],
         "acceleration": [0, 0, 0],
         "orientation": R.from_euler("xyz", [0, 0, 0], degrees=True).as_quat()},
        {"time": 2.0, "position": [1, 2, 0], "velocity": [0, 0, 0],
         "acceleration": [0, 0, 0],
         "orientation": R.from_euler("xyz", [0, 45, 90], degrees=True).as_quat()},
    ]

    interpolator = TrajectoryInterpolator(waypoints)
    ts = np.linspace(0, 2, 100)

    positions = []
    velocities = []
    accelerations = []
    orientations = []
    angular_velocities = []
    angular_accelerations = []

    for t in ts:
        x_r, v_r, a_r = interpolator.interpolate(t)
        positions.append(x_r[:3])
        orientations.append(R.from_quat(x_r[3:]).as_euler('xyz', degrees=True))
        velocities.append(v_r[:3])
        angular_velocities.append(v_r[3:])
        accelerations.append(a_r[:3])
        angular_accelerations.append(a_r[3:])

    positions = np.array(positions)
    orientations = np.array(orientations)
    velocities = np.array(velocities)
    angular_velocities = np.array(angular_velocities)
    accelerations = np.array(accelerations)
    angular_accelerations = np.array(angular_accelerations)

    fig, axs = plt.subplots(6, 1, figsize=(10, 16), sharex=True)

    axs[0].plot(ts, positions)
    axs[0].set_ylabel("Position [m]")
    axs[0].legend(['x', 'y', 'z'])

    axs[1].plot(ts, velocities)
    axs[1].set_ylabel("Velocity [m/s]")
    axs[1].legend(['vx', 'vy', 'vz'])

    axs[2].plot(ts, accelerations)
    axs[2].set_ylabel("Acceleration [m/s²]")
    axs[2].legend(['ax', 'ay', 'az'])

    axs[3].plot(ts, orientations)
    axs[3].set_ylabel("Orientation [deg]")
    axs[3].legend(['roll', 'pitch', 'yaw'])

    axs[4].plot(ts, angular_velocities)
    axs[4].set_ylabel("Ang Vel [rad/s]")
    axs[4].legend(['wx', 'wy', 'wz'])

    axs[5].plot(ts, angular_accelerations)
    axs[5].set_ylabel("Ang Acc [rad/s²]")
    axs[5].set_xlabel("Time [s]")
    axs[5].legend(['awx', 'awy', 'awz'])

    plt.tight_layout()
    plt.show()
