# 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='line', show_animation=True, save_angle_a='angle_A.txt',
                 save_angle_b='angle_B.txt', **kwargs):
    """
    主函数：生成轨迹、计算逆解，并将 theta1 和 theta4 分别保存为两个 txt 文件，逗号分隔
    参数:
        - trajectory_type: 轨迹类型 ('circle', 'line', 'ellipse', 'square', 'triangle')
        - show_animation: 是否显示动画
        - save_angle_a: 保存 theta1（A角）的文件名，设为 None 不保存
        - save_angle_b: 保存 theta4（B角）的文件名，设为 None 不保存
        - **kwargs: 传递给轨迹生成函数的参数
    """
    # 生成轨迹
    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))
        )
    elif trajectory_type == 'ellipse':
        x_list, y_list = generate_ellipse(
            center=kwargs.get('center', (100, 200)),
            rx=kwargs.get('rx', 50),
            ry=kwargs.get('ry', 25)
        )
    elif trajectory_type == 'square':
        x_list, y_list = generate_square(
            side=kwargs.get('side', 60),
            start_point=kwargs.get('start_point', (100, 200))
        )
    elif trajectory_type == 'triangle':
        x_list, y_list = generate_triangle(
            base_length=kwargs.get('base_length', 100),
            height=kwargs.get('height', 80),
            base_center=kwargs.get('base_center', (100, 200))
        )
    else:
        raise ValueError(f"不支持的轨迹类型: {trajectory_type}")

    # 存储角度值的列表
    angle_A_list = []  # theta1 (弧度或角度)
    angle_B_list = []  # theta4 (弧度或角度)

    # 计算每个点的逆运动学并存储角度（以角度制保存）
    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)
            angle_A_list.append(theta1)
            angle_B_list.append(theta4)
            print(f"第 {i} 个点: A角 = {theta1:.2f}°, B角 = {theta4:.2f}°")
        except Exception as e:
            print(f"第 {i} 点逆解失败: {e}")
            # 可选择插入 NaN 或上一个有效值
            angle_A_list.append(np.nan)
            angle_B_list.append(np.nan)

    # ==================== 保存为两个独立的 txt 文件 ====================
    if save_angle_a:
        with open(save_angle_a, 'w') as f:
            # 将所有 A 角度转为字符串，保留2位小数，用逗号连接
            formatted = ",".join([f"{angle:.2f}" for angle in angle_A_list])
            f.write(formatted)
        print(f"\n✅ A角度（theta1）已保存至: {save_angle_a}")

    if save_angle_b:
        with open(save_angle_b, 'w') as f:
            formatted = ",".join([f"{angle:.2f}" for angle in angle_B_list])
            f.write(formatted)
        print(f"✅ B角度（theta4）已保存至: {save_angle_b}")

    # ==================== 可选：显示动画 ====================
    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='line',start=(125, 300), end=(125, 400), show_animation=False)
    #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))