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,
    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)}")

    # --- 设置绘图 ---
    fig, ax = plt.subplots(figsize=(10, 8))
    ax.set_xlim(-200, L0 + 200)  # 根据 L0 调整
    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()

    # --- 创建动画 ---
    ani = FuncAnimation(
        fig, draw_frame,
        frames=len(x_list),
        interval=50,           # 每帧间隔（ms）
        repeat=False,   #循环播放开关
        blit=True
    )

    plt.show()  # <<< 阻塞显示，确保窗口不黑不闪退
    print("动画结束。")


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

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


    # 可选其他轨迹：
    trajectory_config = {
         'func': line_trajectory,
         'params': {'start': (125, 300), 'end': (275, 300)}
     }
    #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】程序结束")