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

# -------------------------
# 参数
# -------------------------
L1 = 100.0   # 连杆1长度
L2 = 100.0   # 连杆2长度

KP, KD = 60.0, 1.0
DEBUG_MODE = True  # 没有电机时开启 DEBUG

# -------------------------
# 逆运动学
# -------------------------
def inverse_2link(x, y, L1, L2):
    D = (x**2 + y**2 - L1**2 - L2**2) / (2 * L1 * L2)
    if abs(D) > 1:
        raise ValueError("目标点超出二连杆工作空间")
    theta2 = np.arctan2(np.sqrt(1 - D**2), D)
    theta1 = np.arctan2(y, x) - np.arctan2(L2*np.sin(theta2), L1 + L2*np.cos(theta2))
    return theta1, theta2

# -------------------------
# 正运动学
# -------------------------
def forward_2link(theta1, theta2, L1, L2):
    x = L1*np.cos(theta1) + L2*np.cos(theta1 + theta2)
    y = L1*np.sin(theta1) + L2*np.sin(theta1 + theta2)
    return x, y

# -------------------------
# 控制两个电机
# -------------------------
def control_two_motors_mit(theta1_rad, theta2_rad):
    if not DEBUG_MODE:
        # motor_control.controlMIT(motor1, KP, KD, theta1_rad, 0.1, 0.0)
        # motor_control.controlMIT(motor2, KP, KD, theta2_rad, 0.1, 0.0)
        pass
    else:
        print(f"[DEBUG] 控制 -> θ1={theta1_rad:.2f}rad, θ2={theta2_rad:.2f}rad")

# -------------------------
# 轨迹生成
# -------------------------
def circle_trajectory(r=60, steps=100):
    angles = np.linspace(0, 2*np.pi, steps)
    x = r * np.cos(angles) + 100
    y = r * np.sin(angles) + 50
    return x, y

def line_trajectory(x0, y0, x1, y1, steps=100):
    x = np.linspace(x0, x1, steps)
    y = np.linspace(y0, y1, steps)
    return x, y

# -------------------------
# 执行轨迹
# -------------------------
def execute_trajectory(x_list, y_list):
    for x, y in zip(x_list, y_list):
        try:
            theta1, theta2 = inverse_2link(x, y, L1, L2)
            control_two_motors_mit(theta1, theta2)
            x_fk, y_fk = forward_2link(theta1, theta2, L1, L2)
            print(f"目标 ({x:.1f},{y:.1f}) -> FK验证 ({x_fk:.1f},{y_fk:.1f})")
        except ValueError as e:
            print(f" {e}")

# -------------------------
# 动画演示
# -------------------------
def run_animation(x_list, y_list):
    fig, ax = plt.subplots()
    ax.set_xlim(-L1-L2, L1+L2)
    ax.set_ylim(-L1-L2, L1+L2)
    line, = ax.plot([], [], 'o-', lw=2)

    def draw_frame(i):
        x, y = x_list[i], y_list[i]
        try:
            theta1, theta2 = inverse_2link(x, y, L1, L2)
            x1 = L1 * np.cos(theta1)
            y1 = L1 * np.sin(theta1)
            line.set_data([0, x1, x], [0, y1, y])
        except:
            line.set_data([], [])
        return line,

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

# -------------------------
# 主程序
# -------------------------
if __name__ == "__main__":
    # 选择轨迹
    x_list, y_list = circle_trajectory(r=60, steps=100)
    # x_list, y_list = line_trajectory(50, 50, 150, 100, steps=100)

    # 执行轨迹
    execute_trajectory(x_list, y_list)

    # 动画演示
    run_animation(x_list, y_list)