# ==================== 五连杆机械臂轨迹控制集成版（S型速度 + 五次插值）====================
# 功能：轨迹生成 + 逆解 + S型速度规划 + 五次多项式平滑控制 + 动画显示 + 关节角度可视化 + FK 验证
# =======================================================================================

import time
import numpy as np
import matplotlib.pyplot as plt
from matplotlib.animation import FuncAnimation
from time import perf_counter  # 高精度计时

# 设置中文字体和负号显示
plt.rcParams['font.sans-serif'] = ['SimHei', 'WenQuanYi Zen Hei', 'FangSong']
plt.rcParams['axes.unicode_minus'] = False

# ------------------------ 调试开关 ------------------------
DEBUG_MODE = True  # <<< 设为 False 控制真实电机

# 导入运动学和轨迹函数（确保路径正确）
try:
    from calculate.ik import inverseF
    from calculate.fk import forwardF
    from calculate.trajectory import (
        circle_trajectory,
        line_trajectory,
        line_trajectory_fix,
        ellipse_trajectory,
        square_trajectory,
        triangle_trajectory
    )
except ImportError as e:
    print(f"【警告】无法导入运动学模块: {e}, 使用 DEBUG_MODE")
    DEBUG_MODE = True

# -------------------- 非 Debug 模式导入硬件库 --------------------
if not DEBUG_MODE:
    try:
        from DM_CAN.DM_CAN import Motor, MotorControl, DM_Motor_Type, Control_Type
        import serial
    except ImportError as e:
        print(f"硬件库导入失败: {e}")
        exit(1)
else:
    print("【DEBUG MODE】: 已启用调试模式，不连接真实硬件。")

# ------------------------ 机械臂参数 ------------------------
L1 = 250  # 左臂长度 (mm)
L2 = 300
L3 = 300
L4 = 250  # 右臂长度
L0 = 250  # 基座右端偏移

# ------------------------ 电机与通信配置 ------------------------
MOTOR1_ID = 0x01
MOTOR2_ID = 0x04
CAN_SERIAL_PORT = '/dev/ttyACM0'
BAUD_RATE = 921600
KP = 50.0
KD = 1.0
DT = 0.001  # 控制周期 1ms

# ------------------------ 全局变量 ------------------------
motor1 = motor2 = motor_control = None
current_theta1 = current_theta4 = 0.0
x_list = y_list = []  # 轨迹点
line = None

# ------------------------ 精确延时 ------------------------
def busy_wait(dt):
    """高精度延时"""
    start = perf_counter()
    while perf_counter() - start < dt:
        pass

# ------------------------ 角度连续性调整 ------------------------
def adjust_angle_continuity(new_angle, prev_angle):
    """
    防止角度跳变（如 3.14 → -3.14），保持连续
    """
    diff = new_angle - prev_angle
    while diff > np.pi:
        diff -= 2 * np.pi
    while diff < -np.pi:
        diff += 2 * np.pi
    return prev_angle + diff

# ==================== 五次多项式插值 ====================
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

# ==================== S型速度规划（5段式）====================
def s_curve_timing(total_time, dt=0.001, max_accel_ratio=2.0):
    """
    生成 S 型速度剖面的时间参数化 s(t) ∈ [0,1]
    使用 5 段：加加速 → 匀加速 → 匀速 → 匀减速 → 减减速
    """
    t_list = np.arange(0, total_time, dt)
    if len(t_list) == 0:
        t_list = np.array([0])

    # 参数：最大加速度设为 4v/T（S型典型）
    v_max = 0.5 * max_accel_ratio * total_time  # 归一化速度
    a_max = 4 * v_max / total_time             # 加速度
    t_jerk = v_max / a_max                     # 加加速时间
    t_acc = 2 * t_jerk                         # 总加速时间
    t_cruise = total_time - 2 * t_acc

    if t_cruise < 0:
        t_acc = total_time / 2
        t_jerk = t_acc / 2
        t_cruise = 0

    s_list = []
    for t in t_list:
        if t <= t_jerk:
            s = (a_max / t_jerk) * t**3 / 6
        elif t <= t_acc - t_jerk:
            t2 = t - t_jerk
            s = (a_max / t_jerk) * t_jerk**3 / 6 + a_max * t_jerk * t2 + 0.5 * a_max * t2**2
        elif t <= t_acc:
            t3 = t - (t_acc - t_jerk)
            s1 = (a_max / t_jerk) * t_jerk**3 / 6
            s2 = a_max * t_jerk * (t_acc - 2*t_jerk) + 0.5 * a_max * (t_acc - 2*t_jerk)**2
            s3 = a_max * t_jerk * t3 - 0.5 * a_max * t3**2
            s = s1 + s2 + s3
        elif t <= t_acc + t_cruise:
            t4 = t - t_acc
            s = s_list[-1] if s_list else 0
            s += v_max * t4
        else:
            t5 = t - (t_acc + t_cruise)
            if t5 <= t_jerk:
                s = 1 - (v_max * (total_time - t) - 0.5 * a_max * t5**2)
            elif t5 <= t_acc - t_jerk:
                s = 1 - (v_max * (total_time - t - t_jerk) - 0.5 * a_max * t_jerk**2)
            else:
                t6 = t5 - (t_acc - t_jerk)
                s = 1 - (a_max / t_jerk) * t6**3 / 6
        s_list.append(s)

    # 归一化到 [0,1]
    s_max = max(s_list) if s_list else 1
    s_normalized = np.array(s_list) / s_max if s_max > 0 else np.zeros_like(s_list)
    return t_list, s_normalized

# ------------------------ 初始化电机 ------------------------
def init_motors():
    global motor1, motor2, motor_control

    if DEBUG_MODE:
        print("【DEBUG】跳过电机初始化")
        motor1 = motor2 = type('Motor', (), {'id': 0})()
        motor_control = type('MotorControl', (), {
            'enable': lambda x: True,
            'disable': lambda x: None,
            'controlMIT': lambda m, kp, kd, pos, vel, torq: None,
            'refresh_motor_status': lambda m: None,
            'switchControlMode': lambda m, mode: None,
            'save_motor_param': lambda m: None
        })()
        return motor1, motor2, motor_control

    try:
        can_serial = serial.Serial(CAN_SERIAL_PORT, BAUD_RATE, timeout=0.5)
        print(f"CAN 串口 {CAN_SERIAL_PORT} 打开成功")
    except Exception as e:
        print(f"无法打开串口: {e}")
        exit(1)

    motor_control = MotorControl(can_serial)
    motor1 = Motor(DM_Motor_Type.DM4310, MOTOR1_ID, 0x11)
    motor2 = Motor(DM_Motor_Type.DM4310, MOTOR2_ID, 0x15)

    motor_control.addMotor(motor1)
    motor_control.addMotor(motor2)
    motor_control.switchControlMode(motor1, Control_Type.MIT)
    motor_control.switchControlMode(motor2, Control_Type.MIT)
    time.sleep(0.1)
    motor_control.save_motor_param(motor1)
    motor_control.save_motor_param(motor2)
    motor_control.enable(motor1)
    motor_control.enable(motor2)
    print("电机已使能。")
    return motor1, motor2, motor_control

# ------------------------ MIT 控制函数 ------------------------
def control_two_motors_mit(theta1, theta4):
    """发送 MIT 控制指令（角度单位：弧度）"""
    global current_theta1, current_theta4

    pos1 = np.degrees(theta1)
    pos4 = np.degrees(theta4)
    vel = 0.0
    torq = 0.0

    if not DEBUG_MODE:
        motor_control.controlMIT(motor1, KP, KD, pos1, vel, torq)
        motor_control.controlMIT(motor2, KP, KD, pos4, vel, torq)
    else:
        print(f"[DEBUG] 控制 -> θ1={pos1:.2f}°, θ4={pos4:.2f}°")

    current_theta1 = theta1
    current_theta4 = theta4

# ------------------------ 平滑轨迹执行函数（带 S 型速度）------------------------
def execute_smooth_trajectory(x_list, y_list,
                              total_time=8.0,
                              kp=KP, kd=KD, dt=DT,
                              read_interval=0.01):
    """
    执行平滑轨迹：使用 S 型速度剖面 + 五次多项式插值
    返回: (time_log, theta1_log, theta4_log, x_fk_log, y_fk_log)
    """
    global motor_control, motor1, motor2

    time_log = []
    theta1_log = []
    theta4_log = []
    x_fk_log = []
    y_fk_log = []

    # FK 参数
    l1, l2, l3, l4, l5 = L1, L2, L3, L4, L0
    omega1 = omega4 = 0.0
    alpha1 = alpha4 = 0.0

    # 获取起始角度
    if not DEBUG_MODE:
        motor_control.refresh_motor_status(motor1)
        motor_control.refresh_motor_status(motor2)
        time.sleep(0.1)
        motor_control.refresh_motor_status(motor1)
        motor_control.refresh_motor_status(motor2)
        start_theta1 = np.radians(motor1.getPosition())
        start_theta4 = np.radians(motor2.getPosition())
    else:
        start_theta1 = start_theta4 = 0.0

    # 生成 S 型时间剖面
    t_profile, s_profile = s_curve_timing(total_time=total_time, dt=dt)
    if len(s_profile) == 0:
        s_profile = np.linspace(0, 1, len(x_list))

    # 插值轨迹点（使用 S 型进度）
    num_s = len(s_profile)
    x_interp = np.interp(s_profile, np.linspace(0, 1, len(x_list)), x_list)
    y_interp = np.interp(s_profile, np.linspace(0, 1, len(y_list)), y_list)

    prev_theta1 = start_theta1
    prev_theta4 = start_theta4

    start_time = perf_counter()

    for i in range(num_s):
        x = x_interp[i]
        y = y_interp[i]

        try:
            target_theta1_raw, target_theta4_raw = inverseF(x, y, L1, L2, L3, L4, L0)
            target_theta1 = float(target_theta1_raw)
            target_theta4 = float(target_theta4_raw)
        except Exception as e:
            print(f"逆解失败: ({x:.2f}, {y:.2f}) -> {e}")
            target_theta1 = prev_theta1
            target_theta4 = prev_theta4
        else:
            target_theta1 = adjust_angle_continuity(target_theta1, prev_theta1)
            target_theta4 = adjust_angle_continuity(target_theta4, prev_theta4)

        # 五次多项式插值（单点瞬时）
        T_seg = dt
        steps = 1  # 每个点只执行一次
        for step in range(steps):
            t_step = step * T_seg / steps
            theta1 = quintic_interpolation(prev_theta1, target_theta1, t_step, T_seg)
            theta4 = quintic_interpolation(prev_theta4, target_theta4, t_step, T_seg)

            control_two_motors_mit(theta1, theta4)

            # FK 验证
            try:
                xc, yc, _, _, _, _, _, _ = forwardF(
                    u1=theta1, u4=theta4, omega1=omega1, omega4=omega4,
                    l1=l1, l2=l2, l3=l3, l4=l4, l5=l5, alpha1=alpha1, alpha4=alpha4
                )
                x_fk_log.append(xc)
                y_fk_log.append(yc)
            except:
                x_fk_log.append(np.nan)
                y_fk_log.append(np.nan)

            # 读取反馈
            now = perf_counter()
            if not DEBUG_MODE and (now - start_time) % read_interval < dt:
                motor_control.refresh_motor_status(motor1)
                motor_control.refresh_motor_status(motor2)

            time_log.append(now - start_time)
            theta1_log.append(np.degrees(theta1))
            theta4_log.append(np.degrees(theta4))

            next_time = start_time + t_profile[i] + dt
            busy_wait(next_time - perf_counter())

        prev_theta1, prev_theta4 = target_theta1, target_theta4

    return time_log, theta1_log, theta4_log, x_fk_log, y_fk_log

# ------------------------ 动画绘制函数 ------------------------
def draw_frame(i):
    x = x_list[i]
    y = y_list[i]
    try:
        theta1, theta4 = inverseF(x, y, L1, L2, L3, L4, L0)
        x2 = L1 * np.cos(theta1)
        y2 = L1 * np.sin(theta1)
        x4 = L0 + L4 * np.cos(theta4)
        y4 = L4 * np.sin(theta4)
        x_coords = [0, x2, x, x4, L0]
        y_coords = [0, y2, y, y4, 0]
        line.set_data(x_coords, y_coords)
    except:
        line.set_data([], [])
    return line,

# ------------------------ 轨迹动画与执行 + 可视化 + FK 验证 ------------------------
def run_trajectory_with_animation(trajectory_func, **params):
    global x_list, y_list, line

    print(f"生成轨迹: {trajectory_func.__name__}")
    x_list, y_list = trajectory_func(**params)
    print(f"轨迹点数: {len(x_list)}")

    # --- 动画 ---
    fig, ax = plt.subplots(figsize=(10, 8))
    ax.set_xlim(-50, L0 + 100)
    ax.set_ylim(0, 500)
    ax.set_aspect('equal')
    ax.grid(True, alpha=0.6)
    ax.set_title(f"五连杆机械臂 - 轨迹: {trajectory_func.__name__}")

    ax.plot(x_list, y_list, 'b--', alpha=0.5, label='目标轨迹')
    line, = ax.plot([], [], 'r-o', lw=3, ms=6, label='机械臂')
    ax.legend()

    ani = FuncAnimation(fig, draw_frame, frames=len(x_list),
                        interval=50, blit=True, repeat=False)
    plt.show()

    # --- 执行并记录 ---
    print("开始执行轨迹（S型速度 + 五次插值）...")
    time_log, theta1_log, theta4_log, x_fk_log, y_fk_log = execute_smooth_trajectory(
        x_list, y_list,
        total_time=8.0,
        kp=50.0, kd=1.0, dt=0.001
    )

    # --- 绘制关节角度 ---
    fig2, (ax1, ax2) = plt.subplots(2, 1, figsize=(10, 6), sharex=True)
    ax1.plot(time_log, theta1_log, 'b-o', markersize=3, linewidth=1.5, label='θ₁ (电机1)')
    ax1.set_ylabel('关节角 θ₁ (°)')
    ax1.set_title('电机关节角度随时间变化')
    ax1.grid(True, alpha=0.5)
    ax1.legend()

    ax2.plot(time_log, theta4_log, 'r-o', markersize=3, linewidth=1.5, label='θ₄ (电机2)')
    ax2.set_xlabel('时间 (秒)')
    ax2.set_ylabel('关节角 θ₄ (°)')
    ax2.grid(True, alpha=0.5)
    ax2.legend()
    plt.tight_layout()
    plt.show()

    # --- FK 验证 ---
    fig3, ax3 = plt.subplots(figsize=(10, 8))
    ax3.plot(x_list, y_list, 'b--', linewidth=2, label='目标轨迹', alpha=0.8)
    ax3.plot(x_fk_log, y_fk_log, 'g-', linewidth=2, marker='o', markersize=3, label='FK 重建轨迹', alpha=0.7)
    ax3.scatter(x_fk_log[0], y_fk_log[0], c='green', s=100, marker='s', label='起点')
    ax3.scatter(x_fk_log[-1], y_fk_log[-1], c='red', s=100, marker='x', label='终点')
    ax3.set_xlabel('X (mm)')
    ax3.set_ylabel('Y (mm)')
    ax3.set_title('FK 验证：目标 vs 实际轨迹')
    ax3.grid(True, alpha=0.5)
    ax3.legend()
    ax3.set_aspect('equal')
    plt.tight_layout()
    plt.show()

    # --- 误差统计 ---
    mask = ~np.isnan(x_fk_log) & ~np.isnan(y_fk_log)
    x_fk_valid = np.array(x_fk_log)[mask]
    y_fk_valid = np.array(y_fk_log)[mask]
    if len(x_fk_valid) == 0:
        print("【警告】FK 数据全为 nan，无法计算误差")
        return

    t_target = np.linspace(0, 1, len(x_list))
    t_fk = np.linspace(0, 1, len(x_fk_valid))
    from scipy.interpolate import interp1d
    try:
        f_x = interp1d(t_fk, x_fk_valid, kind='linear', fill_value='extrapolate')
        f_y = interp1d(t_fk, y_fk_valid, kind='linear', fill_value='extrapolate')
        x_interp = f_x(t_target)
        y_interp = f_y(t_target)
        errors = np.sqrt((x_interp - x_list)**2 + (y_interp - y_list)**2)
        print("\n=== FK 验证误差统计 ===")
        print(f"平均误差: {np.mean(errors):.3f} mm")
        print(f"最大误差: {np.max(errors):.3f} mm")
        print(f"标准差: {np.std(errors):.3f} mm")
    except Exception as e:
        print(f"误差计算失败: {e}")

    # --- 角度范围 ---
    print("\n=== 关节角度范围统计 ===")
    print(f"θ₁ 范围: {min(theta1_log):.2f}° ~ {max(theta1_log):.2f}°")
    print(f"θ₄ 范围: {min(theta4_log):.2f}° ~ {max(theta4_log):.2f}°")

# ------------------------ 主函数 ------------------------
if __name__ == "__main__":
    try:
        init_motors()

        # 选择轨迹
        trajectory_config = {
            'func': line_trajectory_fix,
            'params': {'start': (25, 300), 'end': (200, 300), 'num_points': 100}
        }

        # 示例：圆形轨迹
        # trajectory_config = {
        #     'func': circle_trajectory,
        #     'params': {'center': (100, 300), 'radius': 40, 'num_points': 100}
        # }

        run_trajectory_with_animation(trajectory_config['func'], **trajectory_config['params'])

    except KeyboardInterrupt:
        print("\n【用户中断】")
    except Exception as e:
        print(f"程序异常: {e}")
    finally:
        if not DEBUG_MODE:
            motor_control.disable(motor1)
            motor_control.disable(motor2)
            print("电机已停机。")
        else:
            print("【DEBUG】程序结束")