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

# ------------------------ 调试开关 ------------------------
DEBUG_MODE = False  # True: 模拟电机, False: 真机

# ------------------------ 电机与控制 ------------------------
if DEBUG_MODE:
    class MockMotor:
        def __init__(self):
            self._pos_deg = 45.0
        def getPosition(self):
            return self._pos_deg
        def setPosition(self, deg):
            self._pos_deg = deg
            print(f"[DEBUG] 电机位置 -> {deg:.2f}°")

    class MockControl:
        def enable(self, m): return True
        def disable(self, m): return None
        def controlMIT(self, m, kp, kd, pos, vel, torq):
            deg = np.degrees(pos)
            m.setPosition(deg)

    motor1 = MockMotor()
    motor2 = MockMotor()
    motor_control = MockControl()
else:
    import serial
    from DM_CAN.DM_CAN import Motor, MotorControl, DM_Motor_Type, Control_Type

    CAN_SERIAL_PORT = '/dev/ttyACM0'
    BAUD_RATE = 921600
    MOTOR1_ID = 0x01
    MOTOR2_ID = 0x02

    def init_real_motors():
        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)

        mc = MotorControl(can_serial)
        m1 = Motor(DM_Motor_Type.DM4310, MOTOR1_ID, 0x11)
        m2 = Motor(DM_Motor_Type.DM4310, MOTOR2_ID, 0x12)
        mc.addMotor(m1)
        mc.addMotor(m2)
        mc.switchControlMode(m1, Control_Type.MIT)
        mc.switchControlMode(m2, Control_Type.MIT)
        time.sleep(0.1)
        mc.save_motor_param(m1)
        mc.save_motor_param(m2)
        mc.enable(m1)
        mc.enable(m2)
        print("真实电机已使能。")
        return m1, m2, mc

# ------------------------ 通用函数 ------------------------
def busy_wait(dt):
    start = perf_counter()
    while perf_counter() - start < dt:
        pass

def adjust_angle_continuity(new_angle, prev_angle):
    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 control_two_motors_mit(theta1_rad, theta4_rad, dt=0.01):
    motor_control.controlMIT(motor1, 50, 1, theta1_rad, 0.1, 0.0)
    motor_control.controlMIT(motor2, 50, 1, theta4_rad, 0.1, 0.0)
    busy_wait(dt)

# ------------------------ 移动到零点 ------------------------
def move_to_zero_point(steps=60, dt=0.01):
    theta1_current = np.radians(motor1.getPosition())
    theta4_current = np.radians(motor2.getPosition())
    theta1_list = np.linspace(theta1_current, 0.0, steps)
    theta4_list = np.linspace(theta4_current, 0.0, steps)

    for t1, t4 in zip(theta1_list, theta4_list):
        t1_adj = adjust_angle_continuity(t1, theta1_current)
        t4_adj = adjust_angle_continuity(t4, theta4_current)
        control_two_motors_mit(t1_adj, t4_adj, dt)

    print("已到达零点。")
    print(f"电机1角度: {motor1.getPosition():.2f}°")
    print(f"电机2角度: {motor2.getPosition():.2f}°")

# ------------------------ 示例轨迹动画 ------------------------
def draw_demo_trajectory():
    x_list = np.linspace(0, 200, 50)
    y_list = 100 + 50 * np.sin(np.linspace(0, 2*np.pi, 50))

    fig, ax = plt.subplots(figsize=(8, 6))
    ax.set_xlim(-50, 250)
    ax.set_ylim(0, 250)
    ax.set_aspect('equal')
    ax.grid(True)
    ax.plot(x_list, y_list, 'b--', label='示例轨迹')
    line, = ax.plot([], [], 'ro-', lw=2, label='机械臂')

    def update(i):
        line.set_data([0, x_list[i]], [0, y_list[i]])
        return line,

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

# ------------------------ 主程序 ------------------------
if __name__ == "__main__":
    if not DEBUG_MODE:
        motor1, motor2, motor_control = init_real_motors()
    else:
        print("【DEBUG模式】使用模拟电机")

    print(f"当前电机角度 -> 电机1: {motor1.getPosition():.2f}°, 电机2: {motor2.getPosition():.2f}°")
    print("移动到零点...")
    move_to_zero_point(steps=60, dt=0.01)
    print("显示示例轨迹动画...")
    draw_demo_trajectory()