import rclpy
from rclpy.node import Node
from std_msgs.msg import Float64MultiArray
from sensor_msgs.msg import JointState
import mujoco
import numpy as np
import os
from mujoco import viewer
from ament_index_python.packages import get_package_share_directory
import threading
import time

class MujocoSimNode(Node):
    def __init__(self):
        super().__init__('mujoco_sim_node')
        
        # 配置参数
        self.declare_parameter('model_path', 'models/hive_core_r0.xml')
        self.declare_parameter('simulation_rate', 100.0)  # Hz
        self.declare_parameter('render', True)
        
        model_path = self.get_parameter('model_path').get_parameter_value().string_value
        sim_rate = self.get_parameter('simulation_rate').get_parameter_value().double_value
        self.render = self.get_parameter('render').get_parameter_value().bool_value
        
        # 获取完整模型路径
        package_path = get_package_share_directory('robot_description')
        self.full_model_path = os.path.join(package_path, model_path)
        
        # 初始化MuJoCo模型和数据
        self.model = mujoco.MjModel.from_xml_path(self.full_model_path)
        self.data = mujoco.MjData(self.model)
        
        # 控制指令存储
        self.joint_commands = np.zeros(self.model.nu)
        self.running = True  # 线程运行标志
        
        # 启动Viewer线程（关键修复）
        if self.render:
            self.viewer_thread = threading.Thread(target=self.run_viewer, daemon=True)
            self.viewer_thread.start()
        
        # ROS发布器和订阅器
        self.joint_state_pub = self.create_publisher(JointState, 'joint_states', 10)
        self.cmd_sub = self.create_subscription(Float64MultiArray, 'joint_commands', self.cmd_callback, 10)
        
        self.get_logger().info(f"MuJoCo simulation node started with model: {self.full_model_path}")
        self.get_logger().info(f"Simulation rate: {sim_rate} Hz")

    def run_viewer(self):
        """在独立线程中运行Viewer，避免与ROS线程冲突"""
        # 启动被动Viewer
        with viewer.launch_passive(self.model, self.data) as viewer_handle:
            # 配置相机参数
            viewer_handle.cam.distance = 3.0
            viewer_handle.cam.elevation = -20
            viewer_handle.cam.azimuth = 45
            
            # 保持Viewer运行，直到节点停止
            while self.running:
                self.simulate_step()  # 执行单步仿真
                viewer_handle.sync()  # 同步仿真数据到Viewer
                time.sleep(0.01)  # 降低CPU占用

    def cmd_callback(self, msg):
        """处理接收到的关节控制指令"""
        if len(msg.data) == self.model.nu:
            self.joint_commands = np.array(msg.data)
        else:
            self.get_logger().warn(f"Command dimension mismatch: expected {self.model.nu}, got {len(msg.data)}")

    def simulate_step(self):
        """执行单步仿真并发布状态"""
        if not self.running:
            return
            
        # 设置控制指令
        self.data.ctrl[:] = self.joint_commands
        
        # 运行仿真步
        mujoco.mj_step(self.model, self.data)
        
        # 发布关节状态
        self.publish_joint_states()

    def publish_joint_states(self):
        """发布关节状态到ROS话题"""
        msg = JointState()
        msg.header.stamp = self.get_clock().now().to_msg()
        
        # 获取关节名称
        msg.name = [self.model.joint(i).name for i in range(self.model.njnt)]
        
        # 转换为Python原生float列表
        msg.position = [float(x) for x in self.data.qpos.copy().tolist()]
        msg.velocity = [float(x) for x in self.data.qvel.copy().tolist()]
        msg.effort = [float(x) for x in self.data.qfrc_actuator.copy().tolist()]
        
        self.joint_state_pub.publish(msg)

    def destroy_node(self):
        """清理资源，确保Viewer线程正确退出"""
        self.running = False  # 停止Viewer线程
        if hasattr(self, 'viewer_thread'):
            self.viewer_thread.join()  # 等待线程结束
        super().destroy_node()

def main(args=None):
    rclpy.init(args=args)
    node = MujocoSimNode()
    try:
        rclpy.spin(node)
    except KeyboardInterrupt:
        pass
    finally:
        node.destroy_node()
        if rclpy.ok():
            rclpy.shutdown()

if __name__ == '__main__':
    main()
    