# main1.py

import numpy as np
import matplotlib.pyplot as plt
from matplotlib.animation import FuncAnimation
from caculate.fkik import TwoLinkArm

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

KP, KD = 60.0, 1.0
DEBUG_MODE = True  # 调试模式

# ✅ 创建机械臂实例（关键！）
arm = TwoLinkArm(l1=L1, l2=L2)


# -------------------------
# 控制两个电机（调试用）
# -------------------------
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):
    print("开始执行轨迹...")
    for i, (x, y) in enumerate(zip(x_list, y_list)):
        try:
            theta1, theta2 = arm.ik(x, y)  # ✅ 使用实例调用 ik
            control_two_motors_mit(theta1, theta2)

            _, (x_fk, y_fk) = arm.fk(theta1, theta2)  # 注意：用 _ 忽略第一个返回值 (x1,y1)
            print(f"点 {i:3d}: 目标({x:6.1f},{y:6.1f}) -> FK验证({x_fk:6.1f},{y_fk:6.1f})") # ✅ 使用实例调用 fk
        except ValueError as e:
            print(f"点 {i:3d}: 跳过不可达点 ({x:.1f}, {y:.1f}) -> {e}")
    print("轨迹执行完成。\n")


# -------------------------
# 动画演示
# -------------------------
def run_animation(x_list, y_list):
    print("启动动画...")
    fig, ax = plt.subplots(figsize=(8, 8))
    ax.set_xlim(-200, 200)
    ax.set_ylim(-200, 200)
    ax.set_aspect('equal')
    ax.grid(True, linestyle='--', alpha=0.6)
    ax.set_title("2R 机械臂轨迹跟踪动画")
    ax.set_xlabel("X (mm)")
    ax.set_ylabel("Y (mm)")

    # 绘图元素
    line_arm, = ax.plot([], [], 'o-', lw=2, color='blue', label='机械臂')
    target_point, = ax.plot([], [], 'r*', markersize=10, label='目标点')
    ax.legend()

    def draw_frame(i):
        x_target, y_target = x_list[i], y_list[i]
        try:
            # ✅ 使用实例 arm 调用 ik
            theta1, theta2 = arm.ik(x_target, y_target)
            x1 = L1 * np.cos(theta1)
            y1 = L1 * np.sin(theta1)
            (x1, y1), (x_fk, y_fk) = arm.fk(theta1, theta2)  # 实际末端位置

            # 更新机械臂：基座 -> 关节1 -> 末端
            line_arm.set_data([0, x1, x_fk], [0, y1, y_fk])
        except Exception as e:
            # IK失败，只显示目标
            line_arm.set_data([], [])

        # 始终显示目标点
        target_point.set_data([x_target], [y_target])
        return line_arm, target_point

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

    plt.show()
    return ani  # 防止动画被回收


# -------------------------
# 主程序
# -------------------------
if __name__ == "__main__":
    # 选择轨迹
    x_list, y_list = circle_trajectory(r=40, 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)