# qt_main.py

import numpy as np
import matplotlib.pyplot as plt
from ik import inverseF  # 假设这是你自己的逆运动学函数
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


# 1. 轨迹生成函数：每个轨迹类型独立封装，支持外部传参
# --------------------------------------------------

def generate_circle(center=(100, 300), radius=40):
    from trajectory import circle_trajectory
    return circle_trajectory(center=center, radius=radius)

def generate_line(start=(125, 300), end=(125, 400)):
    from trajectory import line_trajectory
    return line_trajectory(start=start, end=end)

def generate_ellipse(center=(100, 200), rx=50, ry=25):
    from trajectory import ellipse_trajectory
    return ellipse_trajectory(center=center, rx=rx, ry=ry)

def generate_square(side=60, start_point=(100, 200)):
    from trajectory import square_trajectory
    return square_trajectory(side=side, start_point=start_point)

def generate_triangle(base_length=100, height=80, base_center=(100, 200)):
    from trajectory import triangle_trajectory
    return triangle_trajectory(base_length=base_length, height=height, base_center=base_center)


# 4. 主函数：根据轨迹类型调用对应函数并执行
# --------------------------------------------------


def main_of_5dof(trajectory_type='circle', show_animation=True, **kwargs):
    """
    主函数：根据轨迹类型和参数生成轨迹、计算角度、并可视化
    支持传参：
        - circle: center, radius
        - line: start, end
        - ellipse: center, rx, ry
        - square: side, start_point
        - triangle: base_length, height, base_center
        - show_animation: 是否显示动画 (True/False)
    """

    # 根据轨迹类型生成轨迹
    if trajectory_type == 'circle':
        x_list, y_list = generate_circle(
            center=kwargs.get('center', (100, 300)),
            radius=kwargs.get('radius', 40)
        )
    elif trajectory_type == 'line':
        x_list, y_list = generate_line(
            start=kwargs.get('start', (125, 300)),
            end=kwargs.get('end', (125, 400))
        )
    # 其他轨迹类型的处理...

    # 实时计算并输出角度（无论是否显示动画）
    for i in range(len(x_list)):
        x = x_list[i]
        y = y_list[i]
        try:
            theta1, theta4 = inverseF(x, y, L1, L2, L3, L4, L0)
            print(f"第 {i} 个点: theta1 = {np.degrees(theta1):.2f}°, theta4 = {np.degrees(theta4):.2f}°")
        except Exception as e:
            print(f"第 {i} 点计算失败: {e}")

    # 如果需要显示动画
    if show_animation:
        fig, ax = plt.subplots()
        ax.set_xlim(-300, 500)
        ax.set_ylim(0, 500)
        ax.set_aspect('equal')
        ax.grid(True)
        ax.set_title("五连杆末端沿轨迹运动")
        ax.plot(x_list, y_list, 'b--', label='理想轨迹')
        line, = ax.plot([], [], 'r-o', linewidth=2, markersize=6, label='五连杆结构')

        def draw_frame(i):
            x = x_list[i]
            y = y_list[i]

            try:
                theta1, theta4 = inverseF(x, y, L1, L2, L3, L4, L0)
            except Exception as e:
                print(f"第 {i} 帧: 计算失败 -> {e}")
                theta1 = theta4 = None

            if theta1 is None or theta4 is None:
                line.set_data([], [])
                return line,

            # 左右臂坐标计算
            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)
            return line,

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

# 📌 运行主函数
if __name__ == "__main__":
    main_of_5dof(
        trajectory_type='circle',
        center=(150, 250),
        radius=60,
        show_animation=False  # 设置为 False 则不显示动画
    )

    # 示例：其他轨迹使用方式
    # main_of_5dof(trajectory_type='line', start=(0, 0), end=(200, 300), show_animation=False)
    # main_of_5dof(trajectory_type='ellipse', center=(100, 200), rx=80, ry=40)
    # main_of_5dof(trajectory_type='square', side=100, start_point=(100, 200))
    # main_of_5dof(trajectory_type='triangle', base_length=120, height=100, base_center=(100, 200))