import time

import numpy as np
from robot_client.teleoperators.xr import Xr,XrConfig
import matplotlib.pyplot as plt
from matplotlib.animation import FuncAnimation
from scipy.spatial.transform import Rotation
from mpl_toolkits.mplot3d import Axes3D,art3d
import matplotlib.patches as mpatches

def create_cube(center=(0, 0, 0), size=1):
    """创建立方体的顶点和面"""
    # 立方体的8个顶点
    vertices = np.array([
        [-1, -1, -1], [1, -1, -1], [1, 1, -1], [-1, 1, -1],
        [-1, -1, 1], [1, -1, 1], [1, 1, 1], [-1, 1, 1]
    ], dtype=float)
    
    # 缩放并平移
    vertices = vertices * (size / 2)
    vertices = vertices + np.array(center)
    
    # 定义6个面的顶点索引
    faces = [
        [0, 1, 2, 3],  # 底面
        [4, 5, 6, 7],  # 顶面
        [0, 1, 5, 4],  # 前面
        [2, 3, 7, 6],  # 后面
        [1, 2, 6, 5],  # 右面
        [0, 3, 7, 4],  # 左面
    ]
    
    return vertices, faces

def apply_rotation_to_cube(vertices, quaternion):
    """将四元数旋转应用到立方体顶点"""
    q_xyzw = quaternion
    rot = R.from_quat(q_xyzw)
    
    # 计算中心点
    center = np.mean(vertices, axis=0)
    
    # 相对中心的旋转
    rotated_vertices = vertices.copy()
    rotated_vertices -= center
    rotated_vertices = rot.apply(rotated_vertices)
    rotated_vertices += center
    
    return rotated_vertices
def plot_rotating_cube(fig,quat):
    
    
    # 创建立方体
    vertices, faces = create_cube(center=(0, 0, 0), size=2)
    
    # 颜色设置
    face_colors = ['#1f77b4', '#ff7f0e', '#2ca02c', '#d62728', '#9467bd', '#8c564b']
    edge_color = 'black'
    
    # 绘制每个四元数对应的旋转
    # 应用旋转
    rotated_vertices = apply_rotation_to_cube(vertices, quat)
    
    # 绘制3D立方体
    for j, face in enumerate(faces):
        face_verts = rotated_vertices[face]
        ax3d.add_collection3d(
            art3d.Poly3DCollection(
                [face_verts], 
                facecolors=face_colors[j % len(face_colors)], 
                alpha=0.6,
                edgecolors=edge_color,
                linewidths=1
            )
        )
    
    # 绘制坐标轴
    axis_length = 3
    axes = np.array([[axis_length, 0, 0], [0, axis_length, 0], [0, 0, axis_length]])
    rot = R.from_quat(quat if len(quat)==4 else [quat[1], quat[2], quat[3], quat[0]])
    rotated_axes = rot.apply(axes)
    
    colors = ['red', 'green', 'blue']
    labels = ['X', 'Y', 'Z']
    
    for axis, color, label in zip(rotated_axes, colors, labels):
        ax3d.quiver(0, 0, 0, axis[0], axis[1], axis[2], 
                   color=color, arrow_length_ratio=0.1, linewidth=2)
        ax3d.text(axis[0], axis[1], axis[2], label, color=color, fontsize=12, fontweight='bold')
    
    
    return ax3d

# 示例使用
if __name__ == "__main__":
    # 创建一个旋转序列
    from scipy.spatial.transform import Rotation as R
    
    
    """绘制立方体旋转序列"""
    fig = plt.figure(figsize=(12, 8))
    
    # 创建子图布局
    ax3d = fig.add_subplot(121, projection='3d')
    
    # 设置3D坐标轴
    ax3d.set_xlim([-3, 3])
    ax3d.set_ylim([-3, 3])
    ax3d.set_zlim([-3, 3])
    ax3d.set_xlabel('X轴', fontsize=10)
    ax3d.set_ylabel('Y轴', fontsize=10)
    ax3d.set_zlabel('Z轴', fontsize=10)
    ax3d.set_title('旋转立方体', fontsize=12)
    ax3d.view_init(elev=20, azim=45)
    
    fig.suptitle('四元数旋转可视化', fontsize=16, fontweight='bold')
    angle = 1
    rot = R.from_euler('y', angle)
    quat = rot.as_quat()  # (x, y, z, w)
    def update(frame):
        return plot_rotating_cube(fig,ax3d,quat)

    ani = FuncAnimation(fig, update, frames=200, blit=True)#blit是关键

    plt.show()
    xr = Xr(XrConfig())
    xr.connect()
    while True:
        data = xr.get_action()
        plot_rotating_cube([data["left.rx.vel"],data["left.ry.vel"],data["left.rz.vel"],data["left.rw.vel"]])
        time.sleep(200)