# main_animation.py
import numpy as np
import matplotlib.pyplot as plt
from ik import inverseF
from trajectory import circle_trajectory, line_trajectory, ellipse_trajectory, square_trajectory, triangle_trajectory
from matplotlib.animation import FuncAnimation

# 设置中文字体和解决负号显示问题
plt.rcParams['font.sans-serif'] = ['SimHei', 'WenQuanYi Zen Hei', 'FangSong']  # 按优先级选择字体
plt.rcParams['axes.unicode_minus'] = False  # 显示负号 -

# 杆长参数
L1 = 250
L2 = 300
L3 = 300
L4 = 250
L0 = 250

# 设置绘图区域
fig, ax = plt.subplots()
ax.set_xlim(-300, 500)
ax.set_ylim(0, 500)
ax.set_aspect('equal')
ax.grid(True)
ax.set_title("五连杆末端沿轨迹运动")

line, = ax.plot([], [], 'r-o', linewidth=2, markersize=6)
# 选择轨迹类型：
TRAJECTORY_TYPE = 'ellipse'  # 可选: circle, line, ellipse, square, triangle
if TRAJECTORY_TYPE == 'line':
    x_list, y_list = circle_trajectory(center=(100, 300), radius=40)
elif TRAJECTORY_TYPE == 'line':
    x_list, y_list = line_trajectory(start=(125, 300), end=(125, 400))
elif TRAJECTORY_TYPE == 'ellipse':
    x_list, y_list = ellipse_trajectory(center=(100, 200), rx=50, ry=25)
elif TRAJECTORY_TYPE == 'square':
    x_list, y_list = square_trajectory(side=60)
elif TRAJECTORY_TYPE == 'triangle':
    x_list, y_list = triangle_trajectory(base_length=100, height=80)
else:
    raise ValueError("未知的轨迹类型，请选择 circle / line / ellipse / square / triangle")

# 动画函数
def draw_frame(i):
    x = x_list[i]
    y = y_list[i]

    try:
        theta1, theta4 = inverseF(x, y, L1, L2, L3, L4, L0)
        print(theta1)
        print(theta4)
        # 左侧电机臂末端
        x2 = L1 * np.cos(theta1)
        y2 = L1 * np.sin(theta1)
        # 右侧电机臂末端
        x4 = L4 * np.cos(theta4)+L0
        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)
    except Exception as e:
        print(f"第 {i} 帧跳过，错误: {e}")
        line.set_data([], [])

    return line,

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

plt.show()