import time
import numpy as np
import matplotlib.pyplot as plt
from matplotlib.animation import FuncAnimation


# 设置中文字体和解决负号显示问题
plt.rcParams['font.sans-serif'] = ['SimHei', 'WenQuanYi Zen Hei', 'FangSong']  # 按优先级选择字体
plt.rcParams['axes.unicode_minus'] = False  # 显示负号 -
# ------------------------ 调试开关 ------------------------
DEBUG_MODE = True  # <<< True: 只播动画；False: 控制真实电机

# 导入运动学和轨迹函数
from calculate.ik import inverseF
from calculate.trajectory import (
    circle_trajectory,
    line_trajectory,
    line_trajectory_fix,
    ellipse_trajectory,
    square_trajectory,
    triangle_trajectory
)

# -------------------- 只在非 Debug 模式导入 DM_CAN --------------------
if not DEBUG_MODE:
    from DM_CAN.DM_CAN import Motor, MotorControl, DM_Motor_Type, Control_Type, DM_variable
    import serial
    from can_test_motor import motor_control_can_all
else:
    print("【DEBUG MODE】: 已启用调试模式，不连接真实硬件。")

# ------------------------ 机械臂参数 ------------------------
L1 = 250
L2 = 300
L3 = 300
L4 = 250
L0 = 250  # 右侧基座偏移

# ------------------------ DM4310 配置 ------------------------
MOTOR1_ID = 0x01
MOTOR2_ID = 0x04
CAN_SERIAL_PORT = '/dev/ttyACM0'
BAUD_RATE = 921600
KP = 5.0
KD = 1.0

# ------------------------ 轮趣科技电机配置 ------------------------
MOTOR3_ID = 0x05

# 全局变量
motor1 = motor2 = motor3 = motor_control = None
current_theta1 = 0.0
current_theta4 = 0.0
x_list = y_list = []  # 全局轨迹点
line = None  # 动画线对象
ani = None  # 动画对象


# ------------------------ 初始化电机 ------------------------
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
        })()
        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):
    global current_theta1, current_theta4

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

    if not DEBUG_MODE:
        motor_control.controlMIT(motor1, KP, KD, pos1, vel, 0.0)
        motor_control.controlMIT(motor2, KP, KD, pos4, vel, 0.0)
    else:
        print(f"[DEBUG] 模拟控制 -> Motor1: {pos1:.2f}°, Motor2: {pos4:.2f}°")

    current_theta1 = theta1
    current_theta4 = theta4


# ------------------------ 动画帧绘制函数 ------------------------
def draw_frame(i):
    global line, x_list, y_list

    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)
        # 右臂第一杆末端（注意 L0 偏移）
        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)

        # 发送控制指令（在循环中同步）
        control_two_motors_mit(theta1, theta4)

    except Exception as e:
        print(f"第 {i} 帧逆解失败: {e}")
        line.set_data([], [])

    return line,


# ------------------------ 运行轨迹 + 动画 ------------------------
# ------------------------ 运行轨迹 + 动画 + 关节轨迹可视化 ------------------------
def run_trajectory_with_animation(trajectory_func, **params):
    """
    生成轨迹并播放动画，同时发送控制指令，并可视化关节角度变化
    """
    global x_list, y_list, line, ani

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

    # --- 存储关节角度用于后续绘图 ---
    theta1_list = []
    theta4_list = []

    # --- 设置结构动画 ---
    fig, ax = plt.subplots(figsize=(10, 8))
    ax.set_xlim(-200, L0 + 200)
    ax.set_ylim(0, 500)
    ax.set_aspect('equal')
    ax.grid(True, linestyle='--', alpha=0.6)
    ax.set_title(f"五连杆机械臂运动仿真 - 轨迹: {trajectory_func.__name__}", fontsize=14)

    # 绘制目标轨迹（虚线）
    ax.plot(x_list, y_list, 'b--', alpha=0.5, label='目标轨迹')
    line, = ax.plot([], [], 'r-o', linewidth=3, markersize=8, label='机械臂结构')
    ax.legend()

    # --- 动画更新函数（修改：记录角度）---
    def draw_frame(i):
        x = x_list[i]
        y = y_list[i]

        try:
            theta1, theta4 = inverseF(x, y, L1, L2, L3, L4, L0)

            # 记录角度（弧度 → 角度）
            theta1_deg = np.degrees(theta1)
            theta4_deg = np.degrees(theta4)
            theta1_list.append(theta1_deg)
            theta4_list.append(theta4_deg)

            # 左臂第一杆末端
            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)

            # 发送控制指令
            control_two_motors_mit(theta1, theta4)

        except Exception as e:
            print(f"第 {i} 帧逆解失败: {e}")
            line.set_data([], [])
            theta1_list.append(np.nan)
            theta4_list.append(np.nan)

        return line,

    # --- 创建动画 ---
    ani = FuncAnimation(
        fig, draw_frame,
        frames=len(x_list),
        interval=50,
        repeat=False,
        blit=True
    )

    plt.show()  # 显示机械臂动画
    print("动画结束，正在绘制关节角度轨迹...")

    # --- 绘制关节角度曲线 ---
    if len(theta1_list) == 0:
        print("【警告】未记录到任何关节角度数据。")
        return

    time_axis = np.linspace(0, len(theta1_list) * 0.05, len(theta1_list))  # 假设每帧 ~50ms

    fig2, (ax1, ax2) = plt.subplots(2, 1, figsize=(10, 6), sharex=True)

    # 电机1 关节角 θ1
    ax1.plot(time_axis, theta1_list, 'b-o', markersize=4, linewidth=2)
    ax1.set_ylabel('电机1 关节角 θ₁ (°)')
    ax1.set_title('电机关节角度轨迹')
    ax1.grid(True, alpha=0.5)
    ax1.axhline(y=0, color='k', linestyle='--', alpha=0.3)

    # 电机2 关节角 θ4
    ax2.plot(time_axis, theta4_list, 'r-o', markersize=4, linewidth=2)
    ax2.set_xlabel('时间 (秒)')
    ax2.set_ylabel('电机2 关节角 θ₄ (°)')
    ax2.grid(True, alpha=0.5)
    ax2.axhline(y=0, color='k', linestyle='--', alpha=0.3)

    plt.tight_layout()
    plt.show()

    # --- 打印统计信息 ---
    print("\n=== 关节角度统计 ===")
    print(f"θ₁ 范围: {min(theta1_list):.2f}° ~ {max(theta1_list):.2f}°")
    print(f"θ₄ 范围: {min(theta4_list):.2f}° ~ {max(theta4_list):.2f}°")
    print(f"θ₁ 峰峰值: {max(theta1_list) - min(theta1_list):.2f}°")
    print(f"θ₄ 峰峰值: {max(theta4_list) - min(theta4_list):.2f}°")

# ------------------------ 主函数 ------------------------
if __name__ == "__main__":
    # === 在这里设置轨迹类型和参数 ===

    '''trajectory_config = {
        'func': ellipse_trajectory,
        'params': {
            'center': (100, 200),
            'rx': 50,
            'ry': 25,
            'num_points': 100
        }
    }'''


    # 可选其他轨迹：
    trajectory_config = {
         'func': line_trajectory_fix,
         'params': {'start': (25, 300), 'end': (200, 300), 'vx':3.0, 'vy':5.0, 'num_points':20}
     }
    #trajectory_config = {
    #     'func': circle_trajectory,
    #     'params': {'center': (100, 300), 'radius': 40}
    # }

    # 初始化电机
    try:
        motor1, motor2, motor_control = init_motors()

        # 运行带动画的轨迹
        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】程序结束")