# pip install toppra numpy scipy pandas matplotlib
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import toppra as ta
import toppra.algorithm as algo
import toppra.constraint as constraint
import toppra.interpolator as interp

# === 1) 读取 CSV（只取 q1..q7）===
CSV_PATH = "/home/lgv/cmvr/cmvr-es/data/ik_psi_sweep.csv"   #
joint_cols = ["q1","q2","q3","q4","q5","q6","q7"]
df = pd.read_csv(CSV_PATH)
assert all(c in df.columns for c in joint_cols), "CSV 缺少关节列 q1..q7"
waypoints = df[joint_cols].to_numpy()
N, dof = waypoints.shape
assert N >= 2, "至少需要两行关节路点"

# === 2) 构造样条路径（路径参数 s∈[0,1]）===
s_grid = np.linspace(0, 1, N)
path = interp.SplineInterpolator(s_grid, waypoints)

# === 3) 约束（示例值：请改成你的真实上限）===
vmax = np.array([1.5, 1.5, 1.5, 1.5, 1.5, 1.5, 1.5])  # rad/s
amax = np.array([1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0])  # rad/s^2
vlim = np.vstack([-vmax, vmax]).T
alim = np.vstack([-amax, amax]).T
pc_vel = constraint.JointVelocityConstraint(vlim)
pc_acc = constraint.JointAccelerationConstraint(alim)

# === 4) toppra 时间最优轨迹（起止速度=0）===
topp = algo.TOPPRA([pc_vel, pc_acc], path, solver_wrapper="seidel")
traj = topp.compute_trajectory(0.0, 0.0)
T = traj.get_duration()
print(f"Total duration: {T:.4f} s")

# === 5) 采样并绘图 ===
ts = np.linspace(0, T, 300)
q   = traj.eval(ts)    # (300,7)
qd  = traj.evald(ts)
qdd = traj.evaldd(ts)

# 位置
plt.figure()
for i in range(dof):
    plt.plot(ts, q[:, i], label=f"q{i+1}")
plt.xlabel("Time (s)"); plt.ylabel("Position (rad)"); plt.title("Joint Positions"); plt.legend(); plt.tight_layout(); plt.show()

# 速度
plt.figure()
for i in range(dof):
    plt.plot(ts, qd[:, i], label=f"qd{i+1}")
plt.xlabel("Time (s)"); plt.ylabel("Velocity (rad/s)"); plt.title("Joint Velocities"); plt.legend(); plt.tight_layout(); plt.show()

# 加速度
plt.figure()
for i in range(dof):
    plt.plot(ts, qdd[:, i], label=f"qdd{i+1}")
plt.xlabel("Time (s)"); plt.ylabel("Acceleration (rad/s²)"); plt.title("Joint Accelerations"); plt.legend(); plt.tight_layout(); plt.show()