# ==================== 五连杆机械臂轨迹控制（直接角度发送版）====================
# 功能：轨迹生成 + 逆解 + 直接控制 + 动画显示 + 关节角度可视化 + 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 = False   # <<< 设为 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 = 0x02
CAN_SERIAL_PORT = '/dev/ttyACM0'
BAUD_RATE = 921600
KP = 50.0
KD = 1.0
DT = 0.01  # 控制周期 10ms（每点延时）

# ------------------------ 全局变量 ------------------------
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 init_motors():
    global motor1, motor2, motor_control

    if DEBUG_MODE:
        print("【DEBUG】跳过电机初始化")
        # 给个简单的 mock，包含 getPosition 接口（返回度数）
        class MockMotor:
            def __init__(self):
                self._pos_deg = 90.0  # 默认为 90°（即偏移后的值）
            def getPosition(self):
                return self._pos_deg
            def setPosition(self, deg):
                self._pos_deg = deg

        motor1 = MockMotor()
        motor2 = MockMotor()
        class MockControl:
            def enable(self, m): return True
            def disable(self, m): return None
            def controlMIT(self, m, kp, kd, pos, vel, torq):
                # pos 为弧度，记录到mock motor（用于DEBUG可视化一致性）
                try:
                    deg = np.degrees(pos + np.pi/2)  # 因为 motor 的 getPosition 返回带偏移的度值
                    if m is motor1:
                        motor1.setPosition(deg)
                    elif m is motor2:
                        motor2.setPosition(deg)
                except:
                    pass
                return None
            def refresh_motor_status(self, m): return None
            def switchControlMode(self, m, mode): return None
            def save_motor_param(self, m): return None
        motor_control = MockControl()
        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, 0x12)

    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_rad, theta4_rad):
    """
    发送 MIT 控制指令（角度单位：弧度）
    参数: theta1_rad, theta4_rad —— 目标角度（弧度），这里 theta 已去除 ANGLE_OFFSET_RAD
    """
    global current_theta1, current_theta4

    pos1 = theta1_rad
    pos4 = theta4_rad
    vel = 0.1
    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:
        # DEBUG 模式：打印并更新 mock motor（mock control 已在 init 中实现）
        print(f"[DEBUG] 控制 -> θ1={np.degrees(theta1_rad):.2f}°, θ4={np.degrees(theta4_rad):.2f}°")
        motor_control.controlMIT(motor1, KP, KD, pos1, vel, torq)
        motor_control.controlMIT(motor2, KP, KD, pos4, vel, torq)

    current_theta1 = theta1_rad
    current_theta4 = theta4_rad

# ------------------------ 平滑移动到起点（插值过渡）------------------------
def move_to_start_interpolated(start_x, start_y, steps=60, dt=DT):
    """
    从电机当前位置平滑插值过渡到轨迹起点
    返回：time_log_pre, theta1_log_pre, theta4_log_pre, x_fk_log_pre, y_fk_log_pre
    """
    ANGLE_OFFSET_RAD = np.pi / 2

    # 计算目标角度（逆解）
    try:
        theta1_target_raw, theta4_target_raw = inverseF(start_x, start_y, L1, L2, L3, L4, L0)
        theta1_target = float(theta1_target_raw) - ANGLE_OFFSET_RAD
        theta4_target = float(theta4_target_raw) - ANGLE_OFFSET_RAD
    except Exception as e:
        print(f"起点逆解失败: {e}")
        return [], [], [], [], []

    # 获取当前电机角度（转换为去偏移的弧度值）
    if not DEBUG_MODE:
        try:
            motor_control.refresh_motor_status(motor1)
            motor_control.refresh_motor_status(motor2)
            time.sleep(0.05)
            pos1_deg = motor1.getPosition()  # 电机读数通常为度（含偏移）
            pos2_deg = motor2.getPosition()
            current_theta1 = np.radians(pos1_deg) - ANGLE_OFFSET_RAD
            current_theta4 = np.radians(pos2_deg) - ANGLE_OFFSET_RAD
        except Exception as e:
            print(f"读取电机当前位置失败: {e}; 使用默认 0")
            current_theta1 = current_theta4 = 0.0
    else:
        # DEBUG 下使用 mock motor 的 getPosition（init 已设置）
        try:
            pos1_deg = motor1.getPosition()
            pos2_deg = motor2.getPosition()
            current_theta1 = np.radians(pos1_deg) - ANGLE_OFFSET_RAD
            current_theta4 = np.radians(pos2_deg) - ANGLE_OFFSET_RAD
        except:
            current_theta1 = current_theta4 = 0.0

    # 生成插值序列（考虑角度连续性）
    theta1_target = adjust_angle_continuity(theta1_target, current_theta1)
    theta4_target = adjust_angle_continuity(theta4_target, current_theta4)
    theta1_list = np.linspace(current_theta1, theta1_target, steps)
    theta4_list = np.linspace(current_theta4, theta4_target, steps)

    # 日志
    time_log = []
    theta1_log = []
    theta4_log = []
    x_fk_log = []
    y_fk_log = []

    print(f"平滑移动到起点 -> θ1_target={np.degrees(theta1_target):.2f}°, θ4_target={np.degrees(theta4_target):.2f}° (steps={steps})")

    start_t = perf_counter()
    l1, l2, l3, l4, l5 = L1, L2, L3, L4, L0
    omega1 = omega4 = 0.0
    alpha1 = alpha4 = 0.0

    for t1, t4 in zip(theta1_list, theta4_list):
        # 发送控制
        control_two_motors_mit(float(t1), float(t4))

        # FK 验证（使用带偏移角度）
        try:
            xc, yc, *_ = forwardF(
                u1=float(t1 + ANGLE_OFFSET_RAD),
                u4=float(t4 + ANGLE_OFFSET_RAD),
                omega1=omega1, omega4=omega4,
                l1=l1, l2=l2, l3=l3, l4=l4, l5=l5,
                alpha1=alpha1, alpha4=alpha4
            )
        except Exception:
            xc, yc = np.nan, np.nan

        elapsed = perf_counter() - start_t
        time_log.append(elapsed)
        theta1_log.append(np.degrees(float(t1)))  # 记录发送给电机的角度（度）
        theta4_log.append(np.degrees(float(t4)))
        x_fk_log.append(xc)
        y_fk_log.append(yc)

        busy_wait(dt)

    return time_log, theta1_log, theta4_log, x_fk_log, y_fk_log

# ------------------------ 直接轨迹执行函数（无插值）------------------------
from time import perf_counter
import numpy as np


def execute_direct_trajectory(x_list_in, y_list_in, dt=DT):
    """
    直接执行轨迹：对每个 (x,y) 计算逆解并立即发送角度，不做任何插值或速度规划
    返回: (time_log, theta1_log, theta4_log, x_fk_log, y_fk_log)
    """
    global motor_control, motor1, motor2

    # 定义90度对应于弧度的偏移量
    ANGLE_OFFSET_RAD = np.pi / 2  # 即90度对应的弧度值

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

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

    # 获取当前角度（用于连续性调整）
    if not DEBUG_MODE:
        try:
            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)
            current_theta1 = np.radians(motor1.getPosition()) - ANGLE_OFFSET_RAD
            current_theta4 = np.radians(motor2.getPosition()) - ANGLE_OFFSET_RAD
        except Exception:
            # 若读取失败则以前次 global 值或 0 为准
            current_theta1 = current_theta4 = 0.0
    else:
        # DEBUG 下 mock motor 初始已设置为 90deg，减偏移得到 0
        try:
            current_theta1 = np.radians(motor1.getPosition()) - ANGLE_OFFSET_RAD
            current_theta4 = np.radians(motor2.getPosition()) - ANGLE_OFFSET_RAD
        except Exception:
            current_theta1 = current_theta4 = 0.0

    start_time = perf_counter()

    for i, (x, y) in enumerate(zip(x_list_in, y_list_in)):
        try:
            # 计算逆解
            raw_theta1, raw_theta4 = inverseF(x, y, L1, L2, L3, L4, L0)

            # 应用角度偏移
            target_theta1 = float(raw_theta1) - ANGLE_OFFSET_RAD
            target_theta4 = float(raw_theta4) - ANGLE_OFFSET_RAD

            # 调整角度连续性
            target_theta1 = adjust_angle_continuity(target_theta1, current_theta1)
            target_theta4 = adjust_angle_continuity(target_theta4, current_theta4)

        except Exception as e:
            print(f"第 {i} 点逆解失败 ({x:.2f}, {y:.2f}): {e}")
            # 保持上一角度
            target_theta1 = current_theta1
            target_theta4 = current_theta4
        else:
            current_theta1, current_theta4 = target_theta1, target_theta4

        # === 直接发送角度到电机（假设motor api接受弧度值）===
        control_two_motors_mit(target_theta1, target_theta4)

        # === FK 验证实际到达位置 ===
        try:
            xc, yc, *_ = forwardF(
                u1=target_theta1 + ANGLE_OFFSET_RAD,
                u4=target_theta4 + ANGLE_OFFSET_RAD,
                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 Exception:
            x_fk_log.append(np.nan)
            y_fk_log.append(np.nan)

        # 记录日志
        elapsed = perf_counter() - start_time
        time_log.append(elapsed)
        theta1_log.append(np.degrees(target_theta1))  # 将弧度转回角度以便记录（电机实际角度）
        theta4_log.append(np.degrees(target_theta4))

        # 固定延时，控制发送频率
        busy_wait(dt)

    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 Exception:
        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)}")

    if len(x_list) == 0:
        print("轨迹点为空，退出。")
        return

    # --- 平滑移动到轨迹起点（插值过渡），并收集 pre-move 日志 ---
    t_pre, th1_pre, th4_pre, xfk_pre, yfk_pre = move_to_start_interpolated(x_list[0], y_list[0], steps=60, dt=DT)

    # --- 动画 ---
    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("开始执行轨迹（正向）...")
    t_fwd, th1_fwd, th4_fwd, xfk_fwd, yfk_fwd = execute_direct_trajectory(x_list, y_list, dt=DT)

    # --- 执行并记录（反向）---
    print("开始执行轨迹（反向）...")
    t_rev, th1_rev, th4_rev, xfk_rev, yfk_rev = execute_direct_trajectory(x_list[::-1], y_list[::-1], dt=DT)

    # --- 合并三个阶段的数据（仅用于角度时间线显示） ---
    time_offset = 0.0
    all_time = []
    all_th1 = []
    all_th4 = []
    all_xfk = []
    all_yfk = []

    # add pre
    for dt0, a, b, c, d in zip(t_pre, th1_pre, th4_pre, xfk_pre, yfk_pre):
        all_time.append(time_offset + dt0)
        all_th1.append(a)
        all_th4.append(b)
        all_xfk.append(c)
        all_yfk.append(d)
    if len(t_pre) > 0:
        time_offset = all_time[-1]

    # add fwd (shifted)
    for dt0, a, b, c, d in zip(t_fwd, th1_fwd, th4_fwd, xfk_fwd, yfk_fwd):
        all_time.append(time_offset + dt0 + DT)
        all_th1.append(a)
        all_th4.append(b)
        all_xfk.append(c)
        all_yfk.append(d)
    if len(t_fwd) > 0:
        time_offset = all_time[-1]

    # add rev (shifted)
    for dt0, a, b, c, d in zip(t_rev, th1_rev, th4_rev, xfk_rev, yfk_rev):
        all_time.append(time_offset + dt0 + DT)
        all_th1.append(a)
        all_th4.append(b)
        all_xfk.append(c)
        all_yfk.append(d)

    # --- 绘制关节角度（包含预定位） ---
    fig2, (ax1, ax2) = plt.subplots(2, 1, figsize=(10, 6), sharex=True)
    ax1.plot(all_time, all_th1, 'b-o', markersize=3, linewidth=1.0, label='θ₁ (电机1)')
    ax1.set_ylabel('关节角 θ₁ (°)')
    ax1.set_title('电机关节角度随时间变化（含预定位 + 正向 + 反向）')
    ax1.grid(True, alpha=0.5)
    ax1.legend()

    ax2.plot(all_time, all_th4, 'r-o', markersize=3, linewidth=1.0, 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)

    # 正向 FK（对比目标）
    if len(xfk_fwd) > 0:
        ax3.plot(xfk_fwd, yfk_fwd, 'g-', linewidth=2, marker='o', markersize=3, label='FK 重建轨迹（正向）', alpha=0.8)
        # 标注正向起点/终点（若存在有效点）
        valid_idx = [i for i, (xx, yy) in enumerate(zip(xfk_fwd, yfk_fwd)) if not (np.isnan(xx) or np.isnan(yy))]
        if len(valid_idx) > 0:
            ax3.scatter(xfk_fwd[valid_idx[0]], yfk_fwd[valid_idx[0]], c='green', s=100, marker='s', label='正向起点')
            ax3.scatter(xfk_fwd[valid_idx[-1]], yfk_fwd[valid_idx[-1]], c='lime', s=100, marker='x', label='正向终点')

    # 反向 FK（对比反向目标）
    if len(xfk_rev) > 0:
        # xfk_rev 对应的是 x_list[::-1]
        ax3.plot(xfk_rev, yfk_rev, 'm-', linewidth=1.5, marker='x', markersize=3, label='FK 重建轨迹（反向）', alpha=0.7)

    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', adjustable='box')
    plt.tight_layout()
    plt.show()

    # --- 误差统计（仅统计正向段的误差，与目标 x_list/y_list 对比） ---
    mask = ~np.isnan(xfk_fwd) & ~np.isnan(yfk_fwd)
    if np.sum(mask) == 0:
        print("【警告】正向段 FK 数据全为 nan，无法计算误差")
    else:
        try:
            valid_x = np.array(xfk_fwd)[mask]
            valid_y = np.array(yfk_fwd)[mask]
            # 用线性插值重采样 valid 序列到目标轨迹长度
            from scipy.interpolate import interp1d
            t_fk = np.linspace(0, 1, len(valid_x))
            f_x = interp1d(t_fk, valid_x, kind='linear', fill_value='extrapolate')
            f_y = interp1d(t_fk, valid_y, kind='linear', fill_value='extrapolate')
            t_target = np.linspace(0, 1, len(x_list))
            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}")

    # --- 关节角度范围（仍显示整个过程）---
    if len(all_th1) > 0 and len(all_th4) > 0:
        print("\n=== 关节角度范围统计 ===")
        print(f"θ₁ 范围: {min(all_th1):.2f}° ~ {max(all_th1):.2f}°")
        print(f"θ₄ 范围: {min(all_th4):.2f}° ~ {max(all_th4):.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:
            try:
                motor_control.disable(motor1)
                motor_control.disable(motor2)
                print("电机已停机。")
            except Exception:
                pass
        else:
            print("【DEBUG】程序结束")