import math
import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
from matplotlib.widgets import Slider

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

class SpatialTwoLinkArm:
    def __init__(self, l1=1.0, l2=1.0):
        self.l1 = l1
        self.l2 = l2
        self.theta1 = 0.0  # 绕Z轴旋转（偏航，Yaw）
        self.theta2 = 0.0  # 绕Y轴旋转（俯仰，Pitch）

        # 创建3D图形
        self.fig = plt.figure(figsize=(10, 8))
        self.ax = self.fig.add_subplot(111, projection='3d')
        self.ax.set_xlim(-2, 2)
        self.ax.set_ylim(-2, 2)
        self.ax.set_zlim(0, 2)
        self.ax.set_xlabel("X")
        self.ax.set_ylabel("Y")
        self.ax.set_zlabel("Z")
        self.ax.set_title("空间两自由度机械臂（绕Z轴 + 绕Y轴俯仰）")
        self.ax.grid(True)

        # 绘制机械臂
        self.line_link1, = self.ax.plot([], [], [], 'b-', linewidth=4, label='连杆1')
        self.line_link2, = self.ax.plot([], [], [], 'r-', linewidth=4, label='连杆2')
        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], [0], 'ko', markersize=10, label='基座')

        self.ax.legend()

        # 添加角度显示文本
        self.angle_text = self.ax.text2D(
            0.02, 0.95, '', transform=self.ax.transAxes,
            fontsize=10, bbox=dict(boxstyle="round", facecolor="wheat")
        )

        # 添加滑块控制
        self.slider_ax_theta1 = plt.axes([0.2, 0.1, 0.5, 0.03])
        self.slider_ax_theta2 = plt.axes([0.2, 0.05, 0.5, 0.03])

        self.slider_theta1 = Slider(self.slider_ax_theta1, 'θ1(绕Z轴)', -np.pi, np.pi, valinit=0, valfmt="%.2f rad")
        self.slider_theta2 = Slider(self.slider_ax_theta2, 'θ2(俯仰角)', -np.pi/2, np.pi/2, valinit=0, valfmt="%.2f rad")

        self.slider_theta1.on_changed(self.on_slider_change)
        self.slider_theta2.on_changed(self.on_slider_change)

        # 初始绘制
        self.update_plot()

    def rotation_matrix_z(self, theta):
        """绕Z轴旋转矩阵"""
        return np.array([
            [np.cos(theta), -np.sin(theta), 0],
            [np.sin(theta),  np.cos(theta), 0],
            [0,               0,            1]
        ])

    def rotation_matrix_y(self, theta):
        """绕Y轴旋转矩阵"""
        return np.array([
            [np.cos(theta),  0, np.sin(theta)],
            [0,               1,            0],
            [-np.sin(theta), 0, np.cos(theta)]
        ])

    def fk(self, theta1, theta2):
        """正向运动学：返回关节1和末端的位置"""
        # 关节1位置（绕Z旋转后）
        joint1 = np.array([self.l1, 0, 0])
        joint1 = self.rotation_matrix_z(theta1) @ joint1

        # 连杆2的方向：初始沿X轴，先绕Y转θ2，再绕Z转θ1
        link2_dir = np.array([self.l2, 0, 0])
        link2_dir = self.rotation_matrix_z(theta1) @ (self.rotation_matrix_y(theta2) @ link2_dir)

        end_effector = joint1 + link2_dir

        return joint1, end_effector

    def ik(self, x, y, z):
        """逆向运动学：给定末端位置(x,y,z)，求θ1和θ2"""
        # θ1 由 x,y 决定（绕Z轴）
        theta1 = math.atan2(y, x)

        # 在局部坐标系中，x' = sqrt(x^2 + y^2), z' = z
        x_prime = math.sqrt(x*x + y*y)
        z_prime = z

        # 现在是2D平面问题：连杆1长度l1，连杆2长度l2，目标(x_prime, z_prime)
        d_sq = x_prime*x_prime + z_prime*z_prime
        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 = np.clip(cos_theta2_numerator / denominator, -1.0, 1.0)
        theta2 = math.acos(cos_theta2)

        # 修正方向：确保z方向正确
        # 注意：这里假设θ2是向上为正（z>0）
        # 可扩展为elbow_up/down选择

        return theta1, theta2

    def update_plot(self):
        """更新3D图像"""
        joint1, end_effector = self.fk(self.theta1, self.theta2)

        # 确保使用numpy数组以避免shape属性问题
        joint1_np = np.array(joint1)
        end_effector_np = np.array(end_effector)

        # 更新线条
        self.line_link1.set_data_3d([0, joint1_np[0]], [0, joint1_np[1]], [0, joint1_np[2]])
        self.line_link2.set_data_3d([joint1_np[0], end_effector_np[0]], [joint1_np[1], end_effector_np[1]],
                                    [joint1_np[2], end_effector_np[2]])

        self.joint_point.set_data_3d([joint1_np[0]], [joint1_np[1]], [joint1_np[2]])
        self.end_point.set_data_3d([end_effector_np[0]], [end_effector_np[1]], [end_effector_np[2]])

        # 更新角度显示
        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_slider_change(self, val):
        """滑块变化时更新角度"""
        self.theta1 = self.slider_theta1.val
        self.theta2 = self.slider_theta2.val
        self.update_plot()

    def show(self):
        plt.show()


# ========================
# 主程序
# ========================
if __name__ == "__main__":
    print("使用滑块控制空间两自由度机械臂！")
    arm = SpatialTwoLinkArm(l1=1.0, l2=1.0)
    arm.show()