import math
import matplotlib.pyplot as plt

# 设置中文字体（防警告）
plt.rcParams['font.sans-serif'] = ['SimHei', 'Arial Unicode MS', 'DejaVu Sans']
plt.rcParams['axes.unicode_minus'] = False

class TwoLinkArm:
    def __init__(self, l1=1.0, l2=1.0):
        self.l1 = l1
        self.l2 = l2
        self.theta1 = 0.0  # 弧度
        self.theta2 = 0.0  # 弧度
        self.end_x = l1 + l2  # 初始末端位置
        self.end_y = 0.0
        self.joint1 = (l1, 0.0)
        self.joint2 = (l1 + l2, 0.0)

        self.fig, self.ax = plt.subplots(figsize=(8, 8))
        plt.subplots_adjust(bottom=0.2)
        self.ax.set_xlim(-l1 - l2 - 0.5, l1 + l2 + 0.5)
        self.ax.set_ylim(-l1 - l2 - 0.5, l1 + l2 + 0.5)
        self.ax.set_aspect('equal')
        self.ax.grid(True)
        self.ax.set_title("拖动末端执行器（粉色点）来移动机械臂")
        self.ax.set_xlabel("X")
        self.ax.set_ylabel("Y")

        # 绘制机械臂
        self.line_link1, = self.ax.plot([], [], 'b-', linewidth=4, label=f'连杆1 (L1={l1})')
        self.line_link2, = self.ax.plot([], [], 'r-', linewidth=4, label=f'连杆2 (L2={l2})')
        self.joint_point, = self.ax.plot([], [], 'go', markersize=8, label='关节2')
        self.end_point, = self.ax.plot([], [], 'mo', markersize=8, label='末端执行器')
        self.base_point, = self.ax.plot(0, 0, 'ko', markersize=10, label='基座')
        self.angle_text = self.ax.text(
            -l1 - l2 + 0.2,
            l1 + l2 - 0.3,
            '',
            fontsize=10,
            bbox=dict(boxstyle="round", facecolor="wheat")
        )

        self.ax.legend()

        # 拖动相关变量
        self.dragging = False
        self.press_xdata = None
        self.press_ydata = None

        # 连接事件
        self.cid_press = self.fig.canvas.mpl_connect('button_press_event', self.on_press)
        self.cid_release = self.fig.canvas.mpl_connect('button_release_event', self.on_release)
        self.cid_motion = self.fig.canvas.mpl_connect('motion_notify_event', self.on_motion)

        # 初始绘制
        self.update_plot()

    def fk(self, theta1, theta2):
        """正向运动学"""
        x1 = self.l1 * math.cos(theta1)
        y1 = self.l1 * math.sin(theta1)
        x2 = x1 + self.l2 * math.cos(theta1 + theta2)
        y2 = y1 + self.l2 * math.sin(theta1 + theta2)
        return (x1, y1), (x2, y2)

    def ik(self, x, y, elbow_up=True):
        """逆向运动学"""
        d_sq = x*x + y*y
        cos_theta2_numerator = d_sq - self.l1**2 - self.l2**2
        denominator = 2 * self.l1 * self.l2

        if abs(cos_theta2_numerator) > abs(denominator):
            return None  # 不可达

        cos_theta2 = max(-1, min(1, cos_theta2_numerator / denominator))
        theta2 = math.acos(cos_theta2) if elbow_up else -math.acos(cos_theta2)

        k1 = self.l1 + self.l2 * math.cos(theta2)
        k2 = self.l2 * math.sin(theta2)
        theta1 = math.atan2(y, x) - math.atan2(k2, k1)

        return theta1, theta2

    def update_plot(self):
        """根据当前末端位置计算 IK 并更新图像"""
        solution = self.ik(self.end_x, self.end_y, elbow_up=True)  # 默认肘上解
        if solution is not None:
            self.theta1, self.theta2 = solution
            (x1, y1), (x2, y2) = self.fk(self.theta1, self.theta2)
            self.joint1 = (x1, y1)
            self.joint2 = (x2, y2)
        else:
            # 如果不可达，保持上一个有效姿态，但末端显示为红色
            x2, y2 = self.end_x, self.end_y
            self.end_point.set_color('red')

        # 更新线条和点
        self.line_link1.set_data([0, self.joint1[0]], [0, self.joint1[1]])
        self.line_link2.set_data([self.joint1[0], self.joint2[0]], [self.joint1[1], self.joint2[1]])
        self.joint_point.set_data(self.joint1[0], self.joint1[1])
        self.end_point.set_data(self.joint2[0], self.joint2[1])
        self.end_point.set_color('m')  # 恢复正常颜色

        # 显示角度（现在是弧度）
        angle_str = f"θ1 = {self.theta1:+.3f} rad\nθ2= {self.theta2:+.3f} rad"
        self.angle_text.set_text(angle_str)

        self.fig.canvas.draw_idle()

    def on_press(self, event):
        if event.inaxes != self.ax:
            return
        if self.end_point.contains(event)[0]:  # 点中末端点
            self.dragging = True
            self.press_xdata = event.xdata
            self.press_ydata = event.ydata
        else:
            self.dragging = False

    def on_release(self, event):
        self.dragging = False

    def on_motion(self, event):
        if not self.dragging or event.inaxes != self.ax:
            return
        # 更新末端位置
        dx = event.xdata - self.press_xdata
        dy = event.ydata - self.press_ydata
        self.end_x += dx
        self.end_y += dy
        self.press_xdata = event.xdata
        self.press_ydata = event.ydata
        self.update_plot()

    def show(self):
        plt.show()


# ========================
# 主程序
# ========================
if __name__ == "__main__":
    print("拖动末端执行器（粉色点）来交互式控制两连杆机械臂！")
    arm = TwoLinkArm(l1=1.0, l2=1.0)
    arm.show()