# -*- coding: utf-8 -*-
'''
@File    : mujoco_simulator.py
@Time    : 2025/11/27 15:54:20
@Author  : wty-yy
@Version : 1.0
@Blog    : https://wty-yy.github.io/
@Desc    : Mujoco Simulator for Robogauge
'''
import mujoco
import mujoco.viewer
from dm_control import mjcf

import re
import time
import imageio
import numpy as np
from pathlib import Path
from typing import Literal, List, Optional, Union

from robogauge.utils.logger import logger
from robogauge.utils.helpers import parse_path
from robogauge.utils.math_utils import get_projected_gravity, quat_rotate_inverse
from robogauge.tasks.simulator.mujoco_config import MujocoConfig
from robogauge.tasks.simulator.sim_data import (
    SimData,
    RobotProprioception, JointState, BaseState, IMUState,
    RigidBodyDynamics, GroundContactState, DynamicsState, VisualState
)
from robogauge.tasks.gauge.goal_data import VelocityGoal

class MujocoSimulator:
    def __init__(self, sim_cfg: MujocoConfig):
        self.cfg = sim_cfg
        self.terrain_xmls = None
        self.robot_xml = None
        self.terrain_spawn_pos = None
        self.robot_spawn_height = None
        self.default_dof_pos = None
        self.invert_yaw = None
        self.viewer = None
        self.offscreen_cam = mujoco.MjvCamera()
        self.renderer = None
        self.vid_writer = None
        self.vid_count = 0
        self._pause = True
        self.n_step = 0
        self.sim_time = 0.0
        self.target_pos = None
        self.target_velocity: Optional[VelocityGoal] = None
        self.penetration_reset_count = 0
        self.robot_name_prefix = None
        self.robot_body_ids = None
        self.robot_body_id_set = set()
        self.dynamic_body_ids = None
        self.dynamic_body_names = []
        self.dynamic_body_mass = None
        self.dynamic_body_inertia_local = None
        self.dynamic_body_iquat = None
        self.foot_body_ids = None
        self.default_diagonal_foot_distance = 0.0
        self.prev_body_com_pos = None
        self.prev_body_rot = None
        self.prev_body_com_lin_vel = None
        self.prev_body_ang_vel = None
        self.zmp_world_pos = None
        self.zmp_draw_enabled = False
        self.zmp_draw_size = 0.03
        self.zmp_draw_height_offset = 0.02
        self.zmp_draw_rgba = np.array([1.0, 0.85, 0.1, 1.0], dtype=np.float32)

    def load(
        self,
        terrain_xmls: List[str] = None,
        robot_xml: str = None,
        terrain_spawn_pos: list = None,
        default_dof_pos: list = None,
        invert_yaw: bool = None,
    ):
        """ Load terrain and robot into the simulator, support re-loading. """
        if terrain_xmls is not None:
            self.terrain_xmls = [parse_path(xml) for xml in terrain_xmls]
        if robot_xml is not None:
            self.robot_xml = parse_path(robot_xml)
        if terrain_spawn_pos is not None:
            self.terrain_spawn_pos = terrain_spawn_pos
        if default_dof_pos is not None:
            self.default_dof_pos = default_dof_pos
        if invert_yaw is not None:
            self.invert_yaw = invert_yaw

        terrain_xmls = self.terrain_xmls
        robot_xml = self.robot_xml
        terrain_spawn_pos = self.terrain_spawn_pos
        if terrain_xmls is None or robot_xml is None:
            raise ValueError("Terrain and robot XML paths must be provided.")
        if default_dof_pos is None:
            raise ValueError("Default DOF positions must be provided.")
        
        # Create MJCF models
        robot_mjcf = mjcf.from_path(robot_xml)
        terrain_mjcf = mjcf.from_path(terrain_xmls[0])
        visual_elem = terrain_mjcf.visual
        global_elem = visual_elem.get_children('global')
        global_elem.offwidth = 1920
        global_elem.offheight = 1080

        for path in terrain_xmls[1:]:
            next_terrain = mjcf.from_path(path)
            terrain_mjcf.attach(next_terrain)
        for j in robot_mjcf.find_all('joint'):
            if j.tag == 'freejoint':
                j.remove()
        robot_base = robot_mjcf.find('body', 'base_link')
        if robot_base is not None:
            origin_robot_height = robot_base.pos.copy() if robot_base.pos is not None else None
            robot_base.pos = [0, 0, 0]  # move base_link translation to terrain_spawn_pos
        else:
            raise ValueError("Robot base_link body not found in the robot MJCF model.")
        attachment_frame = terrain_mjcf.attach(robot_mjcf)
        attachment_frame.add('freejoint', name='root')
        if origin_robot_height is not None:
            terrain_spawn_pos = np.array(terrain_spawn_pos) + origin_robot_height
        attachment_frame.pos = terrain_spawn_pos

        self.close_viewer()
        self.close_video_writer()
        self.mj_physics = mjcf.Physics.from_mjcf_model(terrain_mjcf)
        self.mj_model = self.mj_physics.model.ptr
        self.mj_data = self.mj_physics.data.ptr
        self.mj_model.opt.timestep = self.cfg.physics.simulation_dt
        self.sim_dt = self.cfg.physics.simulation_dt
        if self.invert_yaw:
            self.mj_data.qpos[3] = 0.0
            self.mj_data.qpos[6] = 1.0
        self.mj_data.qpos[7:] = default_dof_pos
        self.robot_name_prefix = f'{robot_mjcf.model}/'

        # Domain randomization: base mass
        base_body_name = f'{robot_mjcf.model}/base_link'
        body_id = mujoco.mj_name2id(self.mj_model, mujoco.mjtObj.mjOBJ_BODY, base_body_name)
        assert body_id != -1, f"Body '{base_body_name}' not found in the model."
        if self.cfg.domain_rand.base_mass != 0.0:
            original_mass = self.mj_model.body_mass[body_id]
            new_mass = max(0.01, original_mass + self.cfg.domain_rand.base_mass)
            self.mj_model.body_mass[body_id] = new_mass
            logger.info(f"Randomized base mass: {original_mass:.3f} -> {new_mass:.3f} kg")
        
        # Domain randomization: friction
        if self.cfg.domain_rand.friction != 0.0:
            for i in range(self.mj_model.ngeom):
                # Both change robot friction and terrain friction
                # If one of the two geoms has higher priority, the friction of that geom is used.
                # If both geoms have the save priopirty, the maximum of the two friction is used.
                # (Go2 foot friction is 0.4 and priority is 1, terrain priority is 0 except floor)
                self.mj_model.geom_friction[i][0] = self.cfg.domain_rand.friction
            logger.info(f"Scaled geom friction by factor: {self.cfg.domain_rand.friction:.3f}")
        mujoco.mj_forward(self.mj_model, self.mj_data)

        # Setup offscreen camera
        self.offscreen_cam.type = mujoco.mjtCamera.mjCAMERA_TRACKING
        self.offscreen_cam.trackbodyid = body_id
        self.offscreen_cam.distance = self.cfg.viewer.camera_distance
        self.offscreen_cam.elevation = self.cfg.viewer.camera_elevation
        self.offscreen_cam.azimuth = self.cfg.viewer.camera_azimuth
        self.offscreen_cam.lookat = np.array([0.0, 0.0, 0.0])

        # Setup viewer
        self.headless = self.cfg.viewer.headless
        if not self.headless:
            self.viewer = mujoco.viewer.launch_passive(
                self.mj_model, self.mj_data, key_callback=self.key_callback
            )

            # set viewer.camera to follow robot
            self.viewer.cam.type = mujoco.mjtCamera.mjCAMERA_TRACKING
            self.viewer.cam.trackbodyid = body_id
            self.viewer.cam.distance = self.cfg.viewer.camera_distance
            self.viewer.cam.elevation = self.cfg.viewer.camera_elevation
            self.viewer.cam.azimuth = self.cfg.viewer.camera_azimuth
            self.last_render_time = time.time()
        
        # Setup video writer
        if self.cfg.render.save_video:
            self.renderer = mujoco.Renderer(
                self.mj_model, height=self.cfg.render.height,
                width=self.cfg.render.width
            )

            vid_dir = logger.log_dir / "videos"
            vid_dir.mkdir(parents=True, exist_ok=True)
            vid_path = str(vid_dir / f"sim_video_{self.vid_count:03d}.mp4")
            self.vid_writer = imageio.get_writer(
                vid_path,
                fps=self.cfg.render.video_fps,
            )
            self.vid_frame_skip = int(1 / (self.cfg.render.video_fps * self.sim_dt))
            logger.info(f"Simulation video saved at: {vid_path}")
            self.vid_count += 1

        # Initialize simulation state
        self._pause = False
        self.n_step = 0
        self.sim_time = 0.0
        self.penetration_reset_count = 0
        self.load_dof_limits()
        self.preload_sensors()
        self.preload_dynamics_data()
        self.reset_dynamics_cache()

        # Robot controller placeholders
        self.action = None
        self.p_gains = None
        self.d_gains = None
        self.control_type = None

    def key_callback(self, keycode):
        if keycode == 32:
            self._pause = not self._pause
            logger.info(f"Pause toggled: {self._pause}")

    def step(self) -> SimData:
        """ Simulation step, pause will block thread. """
        while self._pause:
            time.sleep(0.1)
        self.update_torque()
        self.mj_physics.step()

        # Viewer sync
        if self.viewer is not None:
            if self.viewer.is_running():
                self.update_external_rendering(self.viewer, ctype='viewer')
                self.viewer.sync()
                time_untile_next_render = self.cfg.physics.simulation_dt - (
                    time.time() - self.last_render_time
                )
                if time_untile_next_render > 0:
                    time.sleep(time_untile_next_render)
                self.last_render_time = time.time()
            else:
                logger.warning("Viewer closed by user.")
                self.close_viewer()

        # Video recording
        if self.vid_writer is not None and self.n_step % self.vid_frame_skip == 0:
            render_cam = self.viewer.cam if self.viewer is not None else self.offscreen_cam
            # mujoco.mjv_updateCamera(render_cam)
            self.renderer.update_scene(self.mj_data, camera=render_cam)
            self.update_external_rendering(self.renderer, ctype='renderer')
            frame = self.renderer.render()
            self.vid_writer.append_data(frame)

        self.proprio = proprio = RobotProprioception(
            joint=JointState(
                pos=self.get_sensor_data('joint_pos'),
                vel=self.get_sensor_data('joint_vel'),
                torque=self.get_sensor_data('joint_eff'),
                limits=self.dof_limits,
                names=self.dof_names,
            ),
            imu=IMUState(
                pos=self.get_sensor_data('imu_pos'),
                quat=self.get_sensor_data('imu_quat'),
                acc=self.get_sensor_data('imu_acc'),
                lin_vel=self.get_sensor_data('imu_lin_vel'),  # body frame, check direction, go2 is inverted
                ang_vel=self.get_sensor_data('imu_ang_vel'),  # body frame, check direction, go2 is inverted
            ),
            base=BaseState(
                pos=self.mj_data.qpos[:3],      # world frame
                quat=self.mj_data.qpos[3:7],    # world frame
                lin_vel=quat_rotate_inverse(self.mj_data.qpos[3:7], self.mj_data.qvel[:3]),  # body frame
                ang_vel=quat_rotate_inverse(self.mj_data.qpos[3:7], self.mj_data.qvel[3:6]), # body frame
            )
        )
        if self.n_step % int(0.1 / self.sim_dt) == 0:
            logger.log(value=np.mean(proprio.imu.quat - proprio.base.quat), tag="sim/delta_quat", step=self.n_step)
            logger.log(value=np.mean(proprio.imu.ang_vel - proprio.base.ang_vel), tag="sim/delta_ang_vel", step=self.n_step)
            logger.log(value=np.mean(proprio.imu.lin_vel - proprio.base.lin_vel), tag="sim/delta_lin_vel", step=self.n_step)
            logger.log(value=proprio.imu.lin_vel[0], tag="sim/imu_lin_vel_x", step=self.n_step)
            logger.log(value=proprio.imu.lin_vel[1], tag="sim/imu_lin_vel_y", step=self.n_step)
            logger.log(value=proprio.base.lin_vel[0], tag="sim/base_lin_vel_x", step=self.n_step)
            logger.log(value=proprio.base.lin_vel[1], tag="sim/base_lin_vel_y", step=self.n_step)
        if self.n_step == 0:
            self.debug_print_proprio_shapes()
        dynamics = self.get_dynamics_data()
        visual = VisualState()

        sim_data = SimData(
            n_step=self.n_step,
            sim_dt=self.sim_dt,
            sim_time=self.sim_time,
            proprio=proprio,
            dynamics=dynamics,
            visual=visual,
        )

        # input("DEBUG")
        self.n_step += 1
        self.sim_time = self.n_step * self.sim_dt
        self.check_truncation(sim_data)
        return sim_data
    
    def update_external_rendering(self,
            handle: Union[mujoco.viewer.Handle, mujoco.Renderer],
            ctype: Literal['viewer', 'renderer'],
        ):
        """ Update external rendering handle (viewer or renderer). """
        def add_target_sphere(geom_elem):
            mujoco.mjv_initGeom(
                geom_elem,
                type=mujoco.mjtGeom.mjGEOM_SPHERE,
                size=[0.1, 0, 0],
                pos=self.target_pos,
                mat=np.eye(3).flatten(),
                rgba=[1, 0, 0, 1]
            )

        def add_zmp_sphere(geom_elem):
            zmp_pos = np.array(self.zmp_world_pos, dtype=np.float32).copy()
            zmp_pos[2] += self.zmp_draw_height_offset
            mujoco.mjv_initGeom(
                geom_elem,
                type=mujoco.mjtGeom.mjGEOM_SPHERE,
                size=[self.zmp_draw_size, 0, 0],
                pos=zmp_pos,
                mat=np.eye(3).flatten(),
                rgba=np.array(self.zmp_draw_rgba, dtype=np.float32)
            )

        def add_thick_arrow(geom_elem, pos, vec, rgba, scale=0.7):
            vel_norm = np.linalg.norm(vec)
            display_norm = min(vel_norm * scale, 1.0)

            if display_norm < 0.10:
                mujoco.mjv_initGeom(
                    geom_elem,
                    type=mujoco.mjtGeom.mjGEOM_NONE,
                    size=[0,0,0], pos=pos, mat=np.eye(3).flatten(), rgba=[0,0,0,0]
                )
                return

            mat = np.zeros(9)
            target_quat = np.zeros(4)
            vec_normalized = vec / vel_norm
            mujoco.mju_quatZ2Vec(target_quat, vec_normalized)
            mujoco.mju_quat2Mat(mat, target_quat)
            
            mat = mat.reshape(3, 3)
            mat[:, 2] *= display_norm 
            
            mujoco.mjv_initGeom(
                geom_elem,
                type=mujoco.mjtGeom.mjGEOM_ARROW,
                size=[0.02, 0.02, display_norm], # [height, width, length]
                pos=pos,
                mat=mat.flatten(),
                rgba=rgba
            )

        viewer_geom_idx = 0
        if ctype == 'viewer':
            handle.user_scn.ngeom = 0  # reset user scene geometry

        if self.target_pos is not None:
            if ctype == 'viewer':
                add_target_sphere(handle.user_scn.geoms[viewer_geom_idx])
                viewer_geom_idx += 1
            else:
                handle.scene.ngeom += 1
                add_target_sphere(handle.scene.geoms[self.renderer.scene.ngeom - 1])

        if self.zmp_draw_enabled and self.zmp_world_pos is not None:
            if ctype == 'viewer':
                add_zmp_sphere(handle.user_scn.geoms[viewer_geom_idx])
                viewer_geom_idx += 1
            else:
                handle.scene.ngeom += 1
                add_zmp_sphere(handle.scene.geoms[self.renderer.scene.ngeom - 1])
        
        if self.target_velocity is not None:
            base_pos_world = self.mj_data.qpos[:3]
            base_quat = self.mj_data.qpos[3:7]
            
            # rendering arrows start position
            offset_body = np.array([0.0, 0.0, 0.2])
            offset_world = np.zeros(3)
            mujoco.mju_rotVecQuat(offset_world, offset_body, base_quat)
            start_pos = base_pos_world + offset_world

            tgt_vel_body = np.array([self.target_velocity.lin_vel_x, self.target_velocity.lin_vel_y, 0.0])
            
            raw_cur_vel = self.proprio.base.lin_vel
            cur_vel_body = np.array([raw_cur_vel[0], raw_cur_vel[1], 0.0])

            tgt_vel_world = np.zeros(3)
            cur_vel_world = np.zeros(3)
            mujoco.mju_rotVecQuat(tgt_vel_world, tgt_vel_body, base_quat)
            if ctype == 'viewer':
                mujoco.mju_rotVecQuat(cur_vel_world, cur_vel_body, base_quat)
            else:
                mujoco.mju_rotVecQuat(cur_vel_world, cur_vel_body, base_quat)

            COLOR_CMD = [0, 1, 0, 1]   # Green 0x00ff00
            COLOR_REAL = [0, 0, 1, 1]  # Blue  0x0000ff

            if ctype == 'viewer':
                # Cmd Arrow
                add_thick_arrow(handle.user_scn.geoms[viewer_geom_idx], start_pos, tgt_vel_world, COLOR_CMD)
                viewer_geom_idx += 1
                # Real Arrow
                add_thick_arrow(handle.user_scn.geoms[viewer_geom_idx], start_pos, cur_vel_world, COLOR_REAL)
                viewer_geom_idx += 1
            else:
                # Renderer Append
                handle.scene.ngeom += 1
                add_thick_arrow(handle.scene.geoms[handle.scene.ngeom - 1], start_pos, tgt_vel_world, COLOR_CMD)
                handle.scene.ngeom += 1
                add_thick_arrow(handle.scene.geoms[handle.scene.ngeom - 1], start_pos, cur_vel_world, COLOR_REAL)

        if ctype == 'viewer':
            handle.user_scn.ngeom = viewer_geom_idx
    
    def check_penetration(self, threshold: float = -0.02):
        if self.penetration_reset_count >= self.cfg.truncation.penetration_max_reset_num:
            return False, None, None, None
        for i in range(self.mj_data.ncon):
            contact = self.mj_data.contact[i]
            if contact.dist < threshold:
                geom1_name = mujoco.mj_id2name(self.mj_model, mujoco.mjtObj.mjOBJ_GEOM, contact.geom1)
                geom2_name = mujoco.mj_id2name(self.mj_model, mujoco.mjtObj.mjOBJ_GEOM, contact.geom2)
                return True, geom1_name, geom2_name, contact.dist
        return False, None, None, None
    
    def check_truncation(self, sim_data: SimData):
        if self.cfg.truncation.enabled:
            projected_gravity = get_projected_gravity(sim_data.proprio.base.quat)
            if -projected_gravity[2] < np.cos(self.cfg.truncation.projected_gravity_rad):
                raise RuntimeError(f"[Roll Error] Episode truncated due to excessive projected gravity, angle: {np.arccos(-projected_gravity[2]):.3f} rad, projected: {projected_gravity}")

            is_penetrated, geom1, geom2, dist = self.check_penetration(self.cfg.truncation.penetration_threshold)
            if is_penetrated:
                is_err = True
                if self.cfg.truncation.skip_penetration_geoms is not None and (
                    any(skip_geom in geom1.lower() for skip_geom in self.cfg.truncation.skip_penetration_geoms) or
                    any(skip_geom in geom2.lower() for skip_geom in self.cfg.truncation.skip_penetration_geoms)
                ):
                    is_err = False
                if self.cfg.truncation.skip_self_penetration:
                    if geom1.split('/')[0] == geom2.split('/')[0]:
                        is_err = False
                if is_err:
                    self.penetration_reset_count += 1
                    raise RuntimeError(f"[Penetration Error] Episode truncated: Penetration ({geom1} <-> {geom2}), distance: {dist}")
    
    def reset(self):
        """ Reset the simulator to initial state. """
        self.mj_physics.reset()
        if self.invert_yaw:
            self.mj_data.qpos[3] = 0.0
            self.mj_data.qpos[6] = 1.0
        self.mj_data.qpos[7:] = self.default_dof_pos
        mujoco.mj_forward(self.mj_model, self.mj_data)
        self.reset_dynamics_cache()

        self.action = None
        self.zmp_world_pos = None
        if self.viewer is not None:
            self.viewer.sync()

    def setup_action(self,
            action: np.ndarray,
            p_gains: np.ndarray = None,
            d_gains: np.ndarray = None,
            control_type: Literal['P'] = 'P'
        ):
        """ Setup action to the simulator. """
        self.action = action
        self.p_gains = p_gains
        self.d_gains = d_gains
        self.control_type = control_type
    
    def update_torque(self):
        if self.action is None:
            return
        dof_pos = self.proprio.joint.pos
        dof_vel = self.proprio.joint.vel
        if self.control_type == 'P':
            torques = self.p_gains * (self.action - dof_pos) - self.d_gains * dof_vel
        else:
            raise NotImplementedError(f"Control type '{self.control_type}' not implemented.")
        self.mj_data.ctrl[:] = torques
    
    def close_viewer(self):
        """ Close the viewer and video writer. """
        if self.viewer is not None:
            self.viewer.close()
            self.viewer = None
            logger.info("Closing viewer.")
    
    def close_video_writer(self):
        """ Close the video writer if exists. """
        if self.vid_writer is not None:
            self.vid_writer.close()
            self.vid_writer = None
            self.renderer = None
            logger.info("Closing video writer.")
    
    def preload_sensors(self):
        # Preload sensor names
        self.joint_pos_sensor_names = self.find_sensors(tag_name="jointpos")
        self.joint_vel_sensor_names = self.find_sensors(tag_name="jointvel")
        self.joint_eff_sensor_names = self.find_sensors(tag_name="jointactuatorfrc")
        self.imu_quat = self.find_sensors(tag_name="framequat")
        self.imu_ang_vel = self.find_sensors(tag_name="gyro")
        self.imu_acc = self.find_sensors(tag_name="accelerometer")
        self.imu_pos = self.find_sensors(tag_name="framepos")
        self.imu_lin_vel = self.find_sensors(tag_name="framelinvel")
        actuator_names = [mujoco.mj_id2name(self.mj_model, mujoco.mjtObj.mjOBJ_ACTUATOR, i) for i in range(self.mj_model.nu)]
        logger.info(
            f"""\nRobot XML: {self.robot_xml}\n"""
            f"""Robot joint names: {[x.rsplit('/')[-1] for x in self.dof_names]}\n"""
            f"""{'='*20} XML SENSOR NAMES {'='*20}\n"""
            f"""Joint Position Sensors [{len(self.joint_pos_sensor_names)}]: {[x.rsplit('/')[-1] for x in self.joint_pos_sensor_names]}\n"""
            f"""Joint Velocity Sensors [{len(self.joint_vel_sensor_names)}]: {[x.rsplit('/')[-1] for x in self.joint_vel_sensor_names]}\n"""
            f"""Joint Effort Sensors [{len(self.joint_eff_sensor_names)}]: {[x.rsplit('/')[-1] for x in self.joint_eff_sensor_names]}\n"""
            f"""Actuators [{len(actuator_names)}]: {[x.rsplit('/')[-1] for x in actuator_names]}\n"""
            f"""IMU Sensors: Quat{self.imu_quat}, AngVel{self.imu_ang_vel}, Acc{self.imu_acc}, Pos{self.imu_pos}, LinVel{self.imu_lin_vel}\n"""
            f"""!!!Checkout actuators order is consistent with joint sensors!!!\n"""
            f"{'='*58}"
        )

        # Cache sensor indices
        self.sensor_cache = {}
        all_lists = {
            'joint_pos': self.joint_pos_sensor_names,
            'joint_vel': self.joint_vel_sensor_names,
            'joint_eff': self.joint_eff_sensor_names,
            'imu_quat': self.imu_quat,
            'imu_ang_vel': self.imu_ang_vel,
            'imu_acc': self.imu_acc,
            'imu_pos': self.imu_pos,
            'imu_lin_vel': self.imu_lin_vel
        }
        
        for key, name_list in all_lists.items():
            indices = []
            for name in name_list:
                sid = mujoco.mj_name2id(self.mj_model, mujoco.mjtObj.mjOBJ_SENSOR, name)
                if sid == -1: continue
                adr = int(self.mj_model.sensor_adr[sid])
                dim = int(self.mj_model.sensor_dim[sid])
                indices.append((adr, dim))
            self.sensor_cache[key] = indices

    def find_sensors(self, *, pattern: re.Pattern = None, tag_name: str = None) -> list:
        model = self.mj_model
        found = []
        tag_map = {
            "jointpos":         mujoco.mjtSensor.mjSENS_JOINTPOS,
            "jointvel":         mujoco.mjtSensor.mjSENS_JOINTVEL,
            "jointactuatorfrc": mujoco.mjtSensor.mjSENS_JOINTACTFRC,
            "accelerometer":    mujoco.mjtSensor.mjSENS_ACCELEROMETER,
            "gyro":             mujoco.mjtSensor.mjSENS_GYRO,
            "framepos":         mujoco.mjtSensor.mjSENS_FRAMEPOS,
            "framequat":        mujoco.mjtSensor.mjSENS_FRAMEQUAT,
            "framelinvel":      mujoco.mjtSensor.mjSENS_FRAMELINVEL,
            "frameangvel":      mujoco.mjtSensor.mjSENS_FRAMEANGVEL,
        }
        tag_type_id = None
        if tag_name:
            if tag_name not in tag_map:
                logger.warning(f"Unknown tag_name '{tag_name}', ignoring tag filter.")
                return []
            tag_type_id = tag_map[tag_name]

        for i in range(model.nsensor):
            if tag_type_id is not None and model.sensor_type[i] != tag_type_id:
                continue
            name = mujoco.mj_id2name(model, mujoco.mjtObj.mjOBJ_SENSOR, i)
            if pattern is None or name and pattern.search(name):
                found.append(name)
        if len(found) == 0:
            logger.warning(f"No sensors found for pattern='{pattern}' tag_name='{tag_name}'")
        return found

    def get_sensor_data(self, cache_key: str) -> np.ndarray:
        ids = self.sensor_cache.get(cache_key, [])
        if not ids:
            return np.array([])

        data_list = []
        for adr, dim in ids:
            data_list.append(self.mj_data.sensordata[adr:adr+dim])
        return np.concatenate(data_list)

    def update_visuals(self, sim_data: SimData):
        visual = sim_data.visual
        if visual is None:
            self.zmp_world_pos = None
            self.zmp_draw_enabled = False
            return

        self.zmp_world_pos = visual.zmp_world_pos
        self.zmp_draw_enabled = bool(visual.zmp_draw_enabled)
        self.zmp_draw_size = float(visual.zmp_draw_size)
        self.zmp_draw_height_offset = float(visual.zmp_draw_height_offset)
        if visual.zmp_draw_rgba is not None:
            self.zmp_draw_rgba = np.array(visual.zmp_draw_rgba, dtype=np.float32)

    @staticmethod
    def _quat_to_rotmat(quat: np.ndarray) -> np.ndarray:
        quat = np.asarray(quat, dtype=np.float64)
        if quat.ndim == 1:
            quat = quat[None, :]
        quat_norm = np.linalg.norm(quat, axis=1, keepdims=True)
        quat_norm = np.maximum(quat_norm, 1e-12)
        quat = quat / quat_norm
        w, x, y, z = quat[:, 0], quat[:, 1], quat[:, 2], quat[:, 3]
        rot = np.empty((quat.shape[0], 3, 3), dtype=np.float32)
        rot[:, 0, 0] = 1.0 - 2.0 * (y * y + z * z)
        rot[:, 0, 1] = 2.0 * (x * y - z * w)
        rot[:, 0, 2] = 2.0 * (x * z + y * w)
        rot[:, 1, 0] = 2.0 * (x * y + z * w)
        rot[:, 1, 1] = 1.0 - 2.0 * (x * x + z * z)
        rot[:, 1, 2] = 2.0 * (y * z - x * w)
        rot[:, 2, 0] = 2.0 * (x * z - y * w)
        rot[:, 2, 1] = 2.0 * (y * z + x * w)
        rot[:, 2, 2] = 1.0 - 2.0 * (x * x + y * y)
        return rot

    @staticmethod
    def _rotation_delta_to_angular_velocity(
            prev_rot: np.ndarray,
            cur_rot: np.ndarray,
            dt: float
        ) -> np.ndarray:
        ang_vel = np.zeros((cur_rot.shape[0], 3), dtype=np.float32)
        if prev_rot is None or dt <= 0.0:
            return ang_vel

        rot_delta = cur_rot @ np.transpose(prev_rot, (0, 2, 1))
        traces = np.trace(rot_delta, axis1=1, axis2=2)
        cos_theta = np.clip((traces - 1.0) * 0.5, -1.0, 1.0)
        theta = np.arccos(cos_theta)
        skew = np.stack([
            rot_delta[:, 2, 1] - rot_delta[:, 1, 2],
            rot_delta[:, 0, 2] - rot_delta[:, 2, 0],
            rot_delta[:, 1, 0] - rot_delta[:, 0, 1],
        ], axis=1)

        small_mask = theta < 1e-6
        if np.any(small_mask):
            ang_vel[small_mask] = 0.5 * skew[small_mask] / dt

        normal_mask = ~small_mask
        if np.any(normal_mask):
            idx = np.where(normal_mask)[0]
            sin_theta = np.sin(theta[idx])
            singular_mask = np.abs(sin_theta) < 1e-6
            if np.any(singular_mask):
                singular_idx = idx[singular_mask]
                ang_vel[singular_idx] = 0.5 * skew[singular_idx] / dt
            regular_idx = idx[~singular_mask]
            if regular_idx.size > 0:
                axis = skew[regular_idx] / (2.0 * sin_theta[~singular_mask, None])
                ang_vel[regular_idx] = axis * theta[regular_idx, None] / dt
        return ang_vel

    @staticmethod
    def _max_pairwise_xy_distance(positions: np.ndarray) -> float:
        if positions.shape[0] < 2:
            return 0.0
        xy = positions[:, :2]
        diffs = xy[:, None, :] - xy[None, :, :]
        dist = np.linalg.norm(diffs, axis=-1)
        return float(np.max(dist))

    def _get_body_name(self, body_id: int) -> str:
        if body_id == 0:
            return "worldbody"
        name = mujoco.mj_id2name(self.mj_model, mujoco.mjtObj.mjOBJ_BODY, body_id)
        return name if name is not None else f"body_{body_id}"

    def _get_geom_name(self, geom_id: int) -> str:
        name = mujoco.mj_id2name(self.mj_model, mujoco.mjtObj.mjOBJ_GEOM, geom_id)
        return name if name is not None else f"geom_{geom_id}"

    def reset_dynamics_cache(self):
        self.prev_body_com_pos = None
        self.prev_body_rot = None
        self.prev_body_com_lin_vel = None
        self.prev_body_ang_vel = None

    def preload_dynamics_data(self):
        self.robot_body_ids = []
        robot_body_names = []
        for body_id in range(self.mj_model.nbody):
            body_name = mujoco.mj_id2name(self.mj_model, mujoco.mjtObj.mjOBJ_BODY, body_id)
            if body_name is not None and body_name.startswith(self.robot_name_prefix):
                self.robot_body_ids.append(body_id)
                robot_body_names.append(body_name)
        self.robot_body_ids = np.array(self.robot_body_ids, dtype=np.int32)
        self.robot_body_id_set = set(int(body_id) for body_id in self.robot_body_ids.tolist())

        if self.robot_body_ids.size == 0:
            logger.warning("No robot bodies found for dynamics preprocessing.")
            self.dynamic_body_ids = np.array([], dtype=np.int32)
            self.dynamic_body_names = []
            self.dynamic_body_mass = np.zeros((0,), dtype=np.float32)
            self.dynamic_body_inertia_local = np.zeros((0, 3), dtype=np.float32)
            self.dynamic_body_iquat = np.zeros((0, 4), dtype=np.float32)
            self.foot_body_ids = np.array([], dtype=np.int32)
            self.default_diagonal_foot_distance = 0.0
            return

        robot_body_mass = np.array(self.mj_model.body_mass[self.robot_body_ids], dtype=np.float32)
        dynamic_mask = robot_body_mass > 0.0
        self.dynamic_body_ids = self.robot_body_ids[dynamic_mask]
        self.dynamic_body_names = [
            robot_body_names[idx]
            for idx, is_dynamic in enumerate(dynamic_mask.tolist())
            if is_dynamic
        ]
        self.dynamic_body_mass = np.array(self.mj_model.body_mass[self.dynamic_body_ids], dtype=np.float32)
        self.dynamic_body_inertia_local = np.array(self.mj_model.body_inertia[self.dynamic_body_ids], dtype=np.float32)
        if hasattr(self.mj_model, 'body_iquat'):
            self.dynamic_body_iquat = np.array(self.mj_model.body_iquat[self.dynamic_body_ids], dtype=np.float32)
        else:
            self.dynamic_body_iquat = np.tile(np.array([[1.0, 0.0, 0.0, 0.0]], dtype=np.float32), (len(self.dynamic_body_ids), 1))

        self.foot_body_ids = np.array([
            body_id for body_id, body_name in zip(self.robot_body_ids.tolist(), robot_body_names)
            if 'foot' in body_name.rsplit('/', 1)[-1].lower()
        ], dtype=np.int32)
        self.default_diagonal_foot_distance = self.compute_default_diagonal_foot_distance()
        logger.info(
            f"Dynamics preprocessing: robot_bodies={len(self.robot_body_ids)}, "
            f"dynamic_bodies={len(self.dynamic_body_ids)}, foot_bodies={len(self.foot_body_ids)}, "
            f"default_diagonal_foot_distance={self.default_diagonal_foot_distance:.4f}"
        )

    def compute_default_diagonal_foot_distance(self) -> float:
        if self.foot_body_ids is not None and self.foot_body_ids.size >= 2:
            if hasattr(self.mj_data, 'xpos'):
                foot_pos = np.array(self.mj_data.xpos[self.foot_body_ids], dtype=np.float32)
            else:
                foot_pos = np.array(self.mj_data.xipos[self.foot_body_ids], dtype=np.float32)
            return self._max_pairwise_xy_distance(foot_pos)

        contacts = self.get_ground_contact_state()
        if contacts.positions.shape[0] >= 2:
            return self._max_pairwise_xy_distance(contacts.positions)
        return 0.0

    def get_ground_contact_state(self) -> GroundContactState:
        positions = []
        distances = []
        normal_forces = []
        tangent_forces = []
        friction_coefficients = []
        robot_geom_names = []
        other_geom_names = []
        robot_body_names = []
        other_body_names = []

        for i in range(self.mj_data.ncon):
            contact = self.mj_data.contact[i]
            geom1 = int(contact.geom1)
            geom2 = int(contact.geom2)
            body1 = int(self.mj_model.geom_bodyid[geom1])
            body2 = int(self.mj_model.geom_bodyid[geom2])
            is_robot1 = body1 in self.robot_body_id_set
            is_robot2 = body2 in self.robot_body_id_set
            if is_robot1 == is_robot2:
                continue

            robot_geom = geom1 if is_robot1 else geom2
            other_geom = geom2 if is_robot1 else geom1
            robot_body = body1 if is_robot1 else body2
            other_body = body2 if is_robot1 else body1

            contact_force = np.zeros(6, dtype=np.float64)
            mujoco.mj_contactForce(self.mj_model, self.mj_data, i, contact_force)
            max_contact_force = 1e6
            contact_force = np.nan_to_num(
                contact_force,
                nan=0.0, posinf=max_contact_force, neginf=-max_contact_force
            )
            contact_force = np.clip(contact_force, -max_contact_force, max_contact_force)

            friction = np.array(contact.friction, dtype=np.float32).reshape(-1)

            positions.append(np.array(contact.pos, dtype=np.float32))
            distances.append(float(contact.dist))
            normal_forces.append(float(abs(contact_force[0])))
            tangent_forces.append(float(np.linalg.norm(contact_force[1:3])))
            friction_coefficients.append(float(friction[0]) if friction.size > 0 else 0.0)
            robot_geom_names.append(self._get_geom_name(robot_geom))
            other_geom_names.append(self._get_geom_name(other_geom))
            robot_body_names.append(self._get_body_name(robot_body))
            other_body_names.append(self._get_body_name(other_body))

        positions = np.array(positions, dtype=np.float32).reshape(-1, 3)
        distances = np.array(distances, dtype=np.float32)
        normal_forces = np.array(normal_forces, dtype=np.float32)
        tangent_forces = np.array(tangent_forces, dtype=np.float32)
        friction_coefficients = np.array(friction_coefficients, dtype=np.float32)
        return GroundContactState(
            positions=positions,
            distances=distances,
            normal_forces=normal_forces,
            tangent_forces=tangent_forces,
            friction_coefficients=friction_coefficients,
            robot_geom_names=robot_geom_names,
            other_geom_names=other_geom_names,
            robot_body_names=robot_body_names,
            other_body_names=other_body_names,
        )

    def get_dynamics_data(self) -> DynamicsState:
        if self.dynamic_body_ids is None or self.dynamic_body_ids.size == 0:
            empty_body = RigidBodyDynamics(
                names=[],
                mass=np.zeros((0,), dtype=np.float32),
                com_pos=np.zeros((0, 3), dtype=np.float32),
                com_lin_vel=np.zeros((0, 3), dtype=np.float32),
                com_lin_acc=np.zeros((0, 3), dtype=np.float32),
                ang_vel=np.zeros((0, 3), dtype=np.float32),
                ang_acc=np.zeros((0, 3), dtype=np.float32),
                inertia_world=np.zeros((0, 3, 3), dtype=np.float32),
            )
            return DynamicsState(
                gravity=np.array(self.mj_model.opt.gravity, dtype=np.float32),
                rigid_bodies=empty_body,
                contacts=self.get_ground_contact_state(),
                default_diagonal_foot_distance=float(self.default_diagonal_foot_distance),
            )

        com_pos = np.array(self.mj_data.xipos[self.dynamic_body_ids], dtype=np.float32)
        body_rot = np.array(self.mj_data.xmat[self.dynamic_body_ids], dtype=np.float32).reshape(-1, 3, 3)
        com_lin_vel = np.zeros_like(com_pos)
        com_lin_acc = np.zeros_like(com_pos)
        ang_vel = np.zeros_like(com_pos)
        ang_acc = np.zeros_like(com_pos)

        if self.prev_body_com_pos is not None:
            com_lin_vel = (com_pos - self.prev_body_com_pos) / self.sim_dt
            if self.prev_body_com_lin_vel is not None:
                com_lin_acc = (com_lin_vel - self.prev_body_com_lin_vel) / self.sim_dt
        if self.prev_body_rot is not None:
            ang_vel = self._rotation_delta_to_angular_velocity(self.prev_body_rot, body_rot, self.sim_dt)
            if self.prev_body_ang_vel is not None:
                ang_acc = (ang_vel - self.prev_body_ang_vel) / self.sim_dt

        if hasattr(self.mj_data, 'ximat'):
            inertial_rot = np.array(self.mj_data.ximat[self.dynamic_body_ids], dtype=np.float32).reshape(-1, 3, 3)
        else:
            inertial_rot = body_rot @ self._quat_to_rotmat(self.dynamic_body_iquat)
        inertia_local = np.zeros((self.dynamic_body_ids.size, 3, 3), dtype=np.float32)
        inertia_local[:, 0, 0] = self.dynamic_body_inertia_local[:, 0]
        inertia_local[:, 1, 1] = self.dynamic_body_inertia_local[:, 1]
        inertia_local[:, 2, 2] = self.dynamic_body_inertia_local[:, 2]
        inertia_world = inertial_rot @ inertia_local @ np.transpose(inertial_rot, (0, 2, 1))

        self.prev_body_com_pos = com_pos.copy()
        self.prev_body_rot = body_rot.copy()
        self.prev_body_com_lin_vel = com_lin_vel.copy()
        self.prev_body_ang_vel = ang_vel.copy()

        return DynamicsState(
            gravity=np.array(self.mj_model.opt.gravity, dtype=np.float32),
            rigid_bodies=RigidBodyDynamics(
                names=self.dynamic_body_names,
                mass=self.dynamic_body_mass.copy(),
                com_pos=com_pos,
                com_lin_vel=com_lin_vel,
                com_lin_acc=com_lin_acc,
                ang_vel=ang_vel,
                ang_acc=ang_acc,
                inertia_world=inertia_world,
            ),
            contacts=self.get_ground_contact_state(),
            default_diagonal_foot_distance=float(self.default_diagonal_foot_distance),
        )

    def debug_print_proprio_shapes(self):
        """Log shapes (or lengths) of each numpy vector inside a RobotProprioception.

        This helps debug mismatched sensor sizes between robots.
        """
        def _shape(x):
            try:
                arr = np.asarray(x)
                return arr.shape
            except Exception:
                return None

        jp = self.proprio.joint
        bs = self.proprio.base
        imu = self.proprio.imu

        logger.info("Proprioception shapes:")
        logger.info(f"  joint.pos: { _shape(jp.pos) }")
        logger.info(f"  joint.vel: { _shape(jp.vel) }")
        logger.info(f"  joint.torque: { _shape(jp.torque) }")

        logger.info(f"  base.pos: { _shape(bs.pos) }")
        logger.info(f"  base.quat: { _shape(bs.quat) }")
        logger.info(f"  base.vel: { _shape(bs.lin_vel) }")
        logger.info(f"  base.ang_vel: { _shape(bs.ang_vel) }")

        logger.info(f"  imu.quat: { _shape(imu.quat) }")
        logger.info(f"  imu.ang_vel: { _shape(imu.ang_vel) }")
        logger.info(f"  imu.acc: { _shape(imu.acc) }")
        logger.info(f"  imu.pos: { _shape(imu.pos) }")
        logger.info(f"  imu.lin_vel: { _shape(imu.lin_vel) }")

    def load_dof_limits(self):
        self.dof_limits = []
        self.dof_names = []
        for i in range(self.mj_model.njnt):
            name = mujoco.mj_id2name(self.mj_model, mujoco.mjtObj.mjOBJ_JOINT, i)
            jnt_type = self.mj_model.jnt_type[i]
            if jnt_type == mujoco.mjtJoint.mjJNT_FREE:
                continue
            limits = self.mj_model.jnt_range[i]
            self.dof_limits.append(limits)
            self.dof_names.append(name)
        self.dof_limits = np.array(self.dof_limits, np.float32)
