import numpy as np
import matplotlib.pyplot as plt
from matplotlib.widgets import Slider, Button
import matplotlib.pyplot as plt
import matplotlib

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

# ======================= 对称五自由度机器人模型 =======================
class Symmetric5DoFRobot:
    def __init__(self):
        self.L1 = 0.25   # 臂段长度
        self.L2 = 0.3
        self.Y0 = 0.25   # 左右臂起始点在 Y 轴上的偏移
        self.tcp_x = 0.8  # 默认 TCP 目标位置
        self.tcp_y = 0.0
        self.tcp_theta = 0  # 夹具方向角度（弧度）

    def inverse_kinematics(self, x, y, theta):
        """给定末端位置和方向，返回左右臂各关节角度"""
        # 计算左臂角度
        try:
            cos_q2_left = (x**2 + y**2 - self.L1**2 - self.L2**2) / (2 * self.L1 * self.L2)
            sin_q2_left = np.sqrt(1 - cos_q2_left**2)
            q2_left = np.arctan2(sin_q2_left, cos_q2_left)

            k1 = self.L1 + self.L2 * np.cos(q2_left)
            k2 = self.L2 * np.sin(q2_left)
            q1_left = np.arctan2(y, x) - np.arctan2(k2, k1)

        except:
            q1_left = 0
            q2_left = 0

        # 右臂镜像求解
        try:
            cos_q2_right = (x**2 + y**2 - self.L1**2 - self.L2**2) / (2 * self.L1 * self.L2)
            sin_q2_right = -np.sqrt(1 - cos_q2_right**2)
            q2_right = np.arctan2(sin_q2_right, cos_q2_right)

            k1 = self.L1 + self.L2 * np.cos(q2_right)
            k2 = self.L2 * np.sin(q2_right)
            q1_right = np.arctan2(y, x) - np.arctan2(k2, k1)

        except:
            q1_right = 0
            q2_right = 0

        return {
            'left': [q1_left, q2_left],
            'right': [q1_right, q2_right],
            'theta': theta
        }

    def forward_kinematics(self):
        ik = self.inverse_kinematics(self.tcp_x, self.tcp_y, self.tcp_theta)
        θ1, θ2 = ik['left']
        θ4, θ3 = ik['right']

        # 左臂坐标计算
        left_base = np.array([0, -self.Y0])
        j1_left = left_base + np.array([self.L1 * np.cos(θ1), self.L1 * np.sin(θ1)])
        tcp_left = j1_left + np.array([
            self.L2 * np.cos(θ1 + θ2),
            self.L2 * np.sin(θ1 + θ2)
        ])

        # 右臂坐标计算
        right_base = np.array([0, self.Y0])
        j1_right = right_base + np.array([self.L1 * np.cos(θ4), self.L1 * np.sin(θ4)])
        tcp_right = j1_right + np.array([
            self.L2 * np.cos(θ4 + θ3),
            self.L2 * np.sin(θ4 + θ3)
        ])

        # 确保末端相连
        tcp_point = (tcp_left + tcp_right) / 2
        gripper_dir = np.array([np.cos(self.tcp_theta), np.sin(self.tcp_theta)]) * 0.3

        return {
            'left_arm': np.array([left_base, j1_left, tcp_point]),
            'right_arm': np.array([right_base, j1_right, tcp_point]),
            'gripper': [tcp_point, tcp_point + gripper_dir]
        }

# ======================= GUI界面 =======================
class Symmetric5DoF_GUI:
    def __init__(self):
        self.robot = Symmetric5DoFRobot()

        self.fig, self.ax = plt.subplots(figsize=(8, 6))
        plt.subplots_adjust(left=0.1, right=0.9, top=0.85, bottom=0.3)
        self.ax.set_title("非共基座对称五自由度机器人（末端相连 + 夹具方向）")
        self.ax.set_xlim(-1.5, 1.5)
        self.ax.set_ylim(-1.5, 1.5)
        self.ax.set_aspect('equal')

        # 初始化线条
        self.left_line, = self.ax.plot([], [], 'b-o', lw=2, markersize=6, label="左臂")
        self.right_line, = self.ax.plot([], [], 'r-o', lw=2, markersize=6, label="右臂")
        self.gripper_line, = self.ax.plot([], [], 'g-->', lw=2, markersize=4, label="夹具方向")

        plt.legend()

        # 创建滑动条
        self.sliders = []
        self.slider_labels = ['X', 'Y', 'Theta']
        self.slider_values = [self.robot.tcp_x, self.robot.tcp_y, self.robot.tcp_theta]

        y_pos = 0.2
        for i in range(3):
            ax_slider = plt.axes([0.2, y_pos, 0.6, 0.03])
            slider = Slider(ax_slider, self.slider_labels[i], -2, 2, valinit=self.slider_values[i])
            slider.on_changed(lambda val, idx=i: self.update_tcp(idx, val))
            self.sliders.append(slider)
            y_pos -= 0.04

        # 添加重置按钮
        reset_ax = plt.axes([0.8, 0.02, 0.1, 0.04])
        self.reset_button = Button(reset_ax, 'Reset', hovercolor='0.975')
        self.reset_button.on_clicked(self.reset)

    def update_tcp(self, idx, value):
        if idx == 0:
            self.robot.tcp_x = value
        elif idx == 1:
            self.robot.tcp_y = value
        elif idx == 2:
            self.robot.tcp_theta = value
        self.update_plot()

    def update_plot(self):
        kinematics = self.robot.forward_kinematics()
        self.left_line.set_data(kinematics['left_arm'][:, 0], kinematics['left_arm'][:, 1])
        self.right_line.set_data(kinematics['right_arm'][:, 0], kinematics['right_arm'][:, 1])

        gb = kinematics['gripper'][0]
        gt = kinematics['gripper'][1]
        self.gripper_line.set_data([gb[0], gt[0]], [gb[1], gt[1]])

        self.fig.canvas.draw_idle()

    def reset(self, event):
        self.robot.tcp_x = 0.8
        self.robot.tcp_y = 0.0
        self.robot.tcp_theta = 0
        for i, s in enumerate(self.sliders):
            s.set_val([0.8, 0, 0][i])
        self.update_plot()

    def run(self):
        self.update_plot()
        plt.show()

# ======================= 主程序入口 =======================
if __name__ == "__main__":
    gui = Symmetric5DoF_GUI()
    gui.run()