# trajectory_control.py

import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
from matplotlib.animation import FuncAnimation

# 导入写好的 fkik 模块
from caculate.fkik import SpatialTwoLinkArm

# -------------------------
# 参数设置
# -------------------------
L1 = 100.0
L2 = 100.0
KP, KD = 60.0, 1.0
DEBUG_MODE = True

# 创建机械臂实例
arm = SpatialTwoLinkArm(l1=L1, l2=L2)

# -------------------------
# 3D 轨迹生成函数
# -------------------------
def circle_trajectory_3d(center=(100, 50, 30), radius=50, axis='xy', steps=100):
    angles = np.linspace(0, 2*np.pi, steps)
    cx, cy, cz = center
    if axis == 'xy':
        x = cx + radius * np.cos(angles)
        y = cy + radius * np.sin(angles)
        z = cz + np.zeros(steps)
    elif axis == 'yz':
        x = cx + np.zeros(steps)
        y = cy + radius * np.cos(angles)
        z = cz + radius * np.sin(angles)
    elif axis == 'xz':
        x = cx + radius * np.cos(angles)
        y = cy + np.zeros(steps)
        z = cz + radius * np.sin(angles)
    else:
        raise ValueError("axis must be 'xy', 'yz', or 'xz'")
    return x, y, z

def line_trajectory_3d(start=(50, 50, 20), end=(150, 100, 80), steps=100):
    t = np.linspace(0, 1, steps)
    x = start[0] + t * (end[0] - start[0])
    y = start[1] + t * (end[1] - start[1])
    z = start[2] + t * (end[2] - start[2])
    return x, y, z

def helix_trajectory_3d(center=(0,0,0), radius=60, height=100, turns=2, steps=200):
    t = np.linspace(0, turns * 2*np.pi, steps)
    x = center[0] + radius * np.cos(t)
    y = center[1] + radius * np.sin(t)
    z = center[2] + height * t / (turns * 2*np.pi)
    return x, y, z

# -------------------------
# 控制电机（调试用）
# -------------------------
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={np.degrees(theta1_rad):+.2f}°, θ2={np.degrees(theta2_rad):+.2f}°")

# -------------------------
# 执行轨迹（调用 IK 和 FK）
# -------------------------
def execute_trajectory(x_list, y_list, z_list):
    for i, (x, y, z) in enumerate(zip(x_list, y_list, z_list)):
        try:
            # 👉 调用你 fkik.py 中的 IK
            solution = arm.ik(x, y, z)
            if solution is None:
                print(f"跳过不可达点 ({x:.1f}, {y:.1f}, {z:.1f})")
                continue
            theta1, theta2 = solution

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

            # 👉 调用 FK 验证
            _, (x_fk, y_fk, z_fk) = arm.fk(theta1, theta2)
            error = np.sqrt((x - x_fk)**2 + (y - y_fk)**2 + (z - z_fk)**2)
            print(f"点 {i}: 目标({x:.1f},{y:.1f},{z:.1f}) -> FK({x_fk:.1f},{y_fk:.1f},{z_fk:.1f}) | 误差={error:.3f}mm")

        except Exception as e:
            print(f"IK 计算失败: {e}")

# -------------------------
# 3D 动画演示
# -------------------------
def run_animation(x_list, y_list, z_list):
    fig = plt.figure("Trajectory Animation")
    ax = fig.add_subplot(111, projection='3d')
    ax.set_xlim(-L1-L2, L1+L2)
    ax.set_ylim(-L1-L2, L1+L2)
    ax.set_zlim(-L1-L2, L1+L2)
    ax.set_xlabel("X")
    ax.set_ylabel("Y")
    ax.set_zlabel("Z")
    ax.set_title("3D Trajectory Tracking")

    # 机械臂线条
    line_arm, = ax.plot([], [], [], 'o-', lw=2, color='blue')
    target_point, = ax.plot([], [], [], 'r*', markersize=10, label='Target')
    ax.legend()

    def animate(frame):
        x, y, z = x_list[frame], y_list[frame], z_list[frame]
        try:
            sol = arm.ik(x, y, z)
            if sol is None:
                return line_arm, target_point
            theta1, theta2 = sol

            # 更新机械臂位置
            joint1, end_effector = arm.fk(theta1, theta2)
            joint1 = np.array(joint1)
            end_effector = np.array(end_effector)

            # 更新线条
            line_arm.set_data_3d([0, joint1[0], end_effector[0]],
                                 [0, joint1[1], end_effector[1]],
                                 [0, joint1[2], end_effector[2]])
            target_point.set_data_3d([x], [y], [z])
        except:
            pass
        return line_arm, target_point

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

# -------------------------
# 主程序
# -------------------------
if __name__ == "__main__":
    print("使用滑块控制空间两自由度机械臂！")

    # 选择一种轨迹
    x_list, y_list, z_list = circle_trajectory_3d(center=(40, -100, -100), radius=20, axis='xy', steps=100)
    # x_list, y_list, z_list = line_trajectory_3d(start=(60, 30, 10), end=(140, 90, 70), steps=80)
    # x_list, y_list, z_list = helix_trajectory_3d(center=(0,0,0), radius=50, height=80, turns=2, steps=200)

    # 执行轨迹（调用 IK 并打印调试信息）
    execute_trajectory(x_list, y_list, z_list)

    # 播放动画
    run_animation(x_list, y_list, z_list)