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

# 显式设置中文字体和正常显示负号
matplotlib.rcParams['font.sans-serif'] = ['WenQuanYi Micro Hei']
matplotlib.rcParams['axes.unicode_minus'] = False  # 解决负号 '-' 显示为方块的问题

# 测试绘图
plt.figure(figsize=(6, 4))
plt.text(0.5, 0.5, '你好，世界！\nThis is 文泉驿微米黑', fontsize=16, ha='center')
plt.axis('off')
plt.title("中文测试")
plt.show()

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

        # 关节限位 [min, max]（单位：弧度）
        self.joint_limits = {
            'theta1': [-np.pi, np.pi],  # 左肩
            'theta2': [-np.pi, np.pi],  # 左肘
            'theta3': [-np.pi, np.pi],  # 右肘
            'theta4': [-np.pi, np.pi],  # 右肩
            'theta5': [-np.pi, np.pi]       # 夹具方向
        }

        # 初始角度
        self.joint_angles = [0.0, 0.0, 0.0, 0.0, 0.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)

            # 限制在合理范围内
            q1_left = np.clip(q1_left, *self.joint_limits['theta1'])
            q2_left = np.clip(q2_left, *self.joint_limits['theta2'])

        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)

            q1_right = np.clip(q1_right, *self.joint_limits['theta4'])
            q2_right = np.clip(q2_right, *self.joint_limits['theta3'])

        except:
            q1_right = 0
            q2_right = 0

        self.joint_angles = [q1_left, q2_left, q2_right, q1_right, theta]
        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.35)
        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.25
        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

        # 添加自由度标签
        joint_names = ['θ1 (Left Shoulder)', 'θ2 (Left Elbow)',
                       'θ3 (Right Elbow)', 'θ4 (Right Shoulder)', 'θ5 (Gripper)']
        limits = self.robot.joint_limits
        ranges = [
            f"[{-round(np.degrees(v[0]), 1)}, {round(np.degrees(v[1]), 1)}]°"
            for v in limits.values()
        ]

        y_pos = 0.15
        for name, rng in zip(joint_names, ranges):
            plt.figtext(0.15, y_pos, f"{name}: {rng}", fontsize=10)
            y_pos -= 0.02

    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 run(self):
        self.update_plot()
        plt.show()

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