import rclpy
import mujoco
import numpy as np
import time
import os

from mujoco import viewer
from .utils import quaternion_to_euler
from ament_index_python.packages import get_package_share_directory
from .s_trajectory import s_curve_jerk_trajectory
from rclpy.node import Node

class RobotMujocoController(Node):
    def __init__(self):
        super().__init__('robot_mujoco_controller')
        
        # 获取模型路径
        package_path = get_package_share_directory('robot_mujoco_control')
        model_path = os.path.join(package_path,"models","robot.xml")

        # 读取XML内容并替换package路径
        try:
            with open(model_path, 'r') as f:
                xml_content = f.read()
            
            # 替换package://路径为绝对路径
            xml_content = xml_content.replace(
                'package://robot_mujoco_control', 
                package_path
            )
            
            # 从处理后的字符串加载模型
            self.model = mujoco.MjModel.from_xml_string(xml_content)
            self.data = mujoco.MjData(self.model)
            self.get_logger().info(f"成功加载模型: {model_path}")
        except Exception as e:
            self.get_logger().error(f"模型加载失败: {str(e)}")
            raise
        
        # 获取执行器索引
        self.actuator_names = [
            "left_leg1_joint_ctrl", 
            "right_leg1_joint_ctrl",
            "left_leg2_joint_ctrl",
            "right_leg2_joint_ctrl",
            "left_wheel1_joint_ctrl",
            "right_wheel1_joint_ctrl",
            "left_wheel2_joint_ctrl",
            "right_wheel2_joint_ctrl",
            "pitch_body_joint_ctrl",
            "left_shoulder_joint_ctrl",
            "right_shoulder_joint_ctrl",
            "left_arm1_joint_ctrl",
            "left_arm2_joint_ctrl",
            "right_arm1_joint_ctrl",
            "right_arm2_joint_ctrl"
        ]
        
        self.actuator_ids = []
        for name in self.actuator_names:
            aid = mujoco.mj_name2id(self.model, mujoco.mjtObj.mjOBJ_ACTUATOR, name)
            if aid == -1:
                self.get_logger().warn(f"执行器 {name} 未找到")
            self.actuator_ids.append(aid)
        
        # 初始化变量
        self.start_time = None
        self.rotation_count = 0
        self.last_yaw = 0
        self.max_speed_factor = 1.0
        self.init_end_time = 3.0
        
    def start_simulation(self):
        # 使用内置查看器
        with viewer.launch_passive(self.model, self.data) as viewer_handle:
            self.start_time = time.time()

            # 设置相机
            viewer_handle.cam.distance = 3.0
            viewer_handle.cam.elevation = -20
            viewer_handle.cam.azimuth = 45
    
            # 运行模拟
            while viewer_handle.is_running():
                current_time = time.time()
                elapsed_time = current_time - self.start_time
                
                # 初始稳定阶段
                if current_time < self.init_end_time + self.start_time:
                    # 保持静止，让机器人稳定
                    # self.data.ctrl[self.actuator_ids] = [-0.349*4.5, -0.349*4.5, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]
                     self.data.ctrl[self.actuator_ids] = [-0.349*2, -0.349*2, 0.349*2, 0.349*2, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]
                else:

                    ## 位置控制
                    cur_dq = [0.0] * 15
                    des_dq = [0.0] * 15
                    tau = [0.0] * 15
                    
                    all_joint_names = [
                        "left_leg1_joint", 
                        "right_leg1_joint", 
                        "left_leg2_joint", 
                        "right_leg2_joint",
                        "left_wheel1_joint",
                        "right_wheel1_joint",
                        "left_wheel2_joint",
                        "right_wheel2_joint",
                        "pitch_body_joint",
                        "left_shoulder_joint",
                        "right_shoulder_joint",
                        "left_arm1_joint",
                        "left_arm2_joint",
                        "right_arm1_joint",
                        "right_arm2_joint"
                        ]
                    
                    wheel_joint_name = [
                        "left_wheel1_joint",
                        "right_wheel1_joint",
                        "left_wheel2_joint",
                        "right_wheel2_joint"
                    ]
                    
                    for i, name in enumerate(all_joint_names):
                        jid = mujoco.mj_name2id(self.model, mujoco.mjtObj.mjOBJ_JOINT, name)
                        if jid != -1:
                            cur_dq[i] = self.data.qpos[jid]
                    
                

                # 步进模拟
                mujoco.mj_step(self.model, self.data)
                
                # 更新查看器
                viewer_handle.sync()
                
                # 获取机器人方向
                base_quat = self.data.qpos[3:7]
                _, _, yaw = quaternion_to_euler(base_quat)
                
                # 检测完整旋转
                if self.last_yaw > 3.0 and yaw < -3.0:  # 从正π到负π
                    self.rotation_count += 1
                elif self.last_yaw < -3.0 and yaw > 3.0:  # 从负π到正π
                    self.rotation_count += 1
                self.last_yaw = yaw
                
                # 控制模拟速度
                time.sleep(0.001)
        
        self.get_logger().info("模拟结束")

def main(args=None):
    rclpy.init(args=args)
    try:
        simObj = RobotMujocoController()
        simObj.start_simulation()
    except KeyboardInterrupt:
        pass
    except Exception as e:
        print(f"发生错误: {str(e)}")
    finally:
        rclpy.shutdown()

if __name__ == '__main__':
    main()
