import time
import numpy as np
import matplotlib.pyplot as plt
from DM_CAN import *
import serial
from time import perf_counter, sleep

# ================= 初始化电机 =================
Motor1 = Motor(DM_Motor_Type.DM4310, 0x01, 0x11)
serial_device = serial.Serial('COM3', 921600, timeout=0.5)
MotorControl1 = MotorControl(serial_device)
MotorControl1.addMotor(Motor1)
MotorControl1.enable(Motor1)
sleep(1)
# ================= 五次多项式轨迹函数 =================
def quintic_interpolation(q0, qf, t, T):
    """
    五次多项式平滑插值
    q0: 起始位置
    qf: 目标位置
    t: 当前时间（0~T）
    T: 总时间
    返回：平滑位置
    """
    s = t / T
    s = np.clip(s, 0, 1)  # 限制在0~1
    alpha = s**3 * (10 - 15*s + 6*s**2)
    return q0 + (qf - q0) * alpha

# ================= 执行轨迹函数 =================
def execute_trajectory(MotorControl, Motor, waypoints,
                       total_time=6.0, slow_down_time=2.0,
                       kp=50, kd=1.0, dt=0.001, read_interval=0.01,
                       ramp_time=2.0):
    """
    五次多项式轨迹执行函数（均分时间 + 缓启动 + 缓停）
    """
    # 获取当前电机位置
    MotorControl.refresh_motor_status(Motor)
    sleep(0.1)
    MotorControl.refresh_motor_status(Motor)
    current_pos = Motor.getPosition()
    positions = [current_pos] + waypoints

    # ===== 均分时间，不按距离 =====
    num_segments = len(positions) - 1
    seg_times = [total_time / num_segments] * num_segments

    # ===== 轨迹执行 =====
    actual_pos, desired_pos, time_axis = [], [], []
    start_time = perf_counter()
    next_time = start_time
    last_read_time = start_time

    for i in range(num_segments):
        q0 = positions[i]
        qf = positions[i+1]
        T = seg_times[i]
        steps = int(T/dt)
        for j in range(steps):
            t = j*dt
            # 五次多项式插值
            pos_cmd = quintic_interpolation(q0, qf, t, T)
            # 启动缓动
            ramp = 0.5*(1 - np.cos(np.pi*min(t, ramp_time)/ramp_time)) if t < ramp_time else 1.0
            pos_cmd = current_pos + (pos_cmd - current_pos) * ramp
            vel_cmd = 0  # 由五次插值保证平滑

            MotorControl.controlMIT(Motor, kp, kd, pos_cmd, vel_cmd, 0)

            # 低频读取
            now = perf_counter()
            if now - last_read_time >= read_interval:
                MotorControl.refresh_motor_status(Motor)
                actual = Motor.getPosition()
                last_read_time = now
            else:
                actual = np.nan

            actual_pos.append(actual)
            desired_pos.append(pos_cmd)
            time_axis.append(now - start_time)

            # 固定周期
            next_time += dt
            while perf_counter() < next_time:
                pass

        current_pos = qf  # 更新起点用于下一段

    # ===== 缓停段 =====
    final_pos = positions[-1]
    MotorControl.refresh_motor_status(Motor)
    start_pos = Motor.getPosition()
    steps = int(slow_down_time / dt)
    for i in range(steps):
        t = i * dt
        pos_cmd = quintic_interpolation(start_pos, final_pos, t, slow_down_time)
        vel_cmd = 0
        MotorControl.controlMIT(Motor, kp, kd, pos_cmd, vel_cmd, 0)

        now = perf_counter()
        if now - last_read_time >= read_interval:
            MotorControl.refresh_motor_status(Motor)
            actual = Motor.getPosition()
            last_read_time = now
        else:
            actual = np.nan

        actual_pos.append(actual)
        desired_pos.append(pos_cmd)
        time_axis.append(now - start_time)

        next_time += dt
        while perf_counter() < next_time:
            pass

    # ===== 绘图 =====
    actual_clean = np.array(actual_pos)
    desired_clean = np.array(desired_pos)
    mask = ~np.isnan(actual_clean)

    plt.figure()
    plt.plot(np.array(time_axis)[mask], desired_clean[mask], label='Desired')
    plt.plot(np.array(time_axis)[mask], actual_clean[mask], label='Actual', linestyle='--')
    plt.xlabel('Time [s]')
    plt.ylabel('Position [rad]')
    plt.title(f'Quintic Trajectory (Ramp + Smooth Stop)')
    plt.legend()
    plt.grid(True)
    plt.show()


# ================= 执行轨迹 =================
target_angles = [ 0.0,  0.3, 0.6, 0.88, 0.9, 1.5,
                  0.5, 0.2,-0.2,-0.8]
execute_trajectory(MotorControl1, Motor1, target_angles,
                   total_time=10.0, slow_down_time=1.0,
                   kp=50, kd=1, dt=0.001, read_interval=0.005,
                   ramp_time=1.0)

# 关闭
MotorControl1.disable(Motor1)
serial_device.close()