from __future__ import annotations

import warnings
from dataclasses import dataclass, field
from pathlib import Path
from typing import TYPE_CHECKING, Callable, Sequence

import mujoco
import mujoco_warp as mjwarp
import numpy as np
import torch

from mjlab import actuator
from mjlab.actuator import BuiltinActuatorGroup, FusedActuatorGroup
from mjlab.actuator.actuator import TransmissionType
from mjlab.actuator.xml_actuator import XmlActuator
from mjlab.entity.data import EntityData
from mjlab.utils import spec_config as spec_cfg
from mjlab.utils.lab_api.string import resolve_matching_names
from mjlab.utils.mujoco import dof_width, qpos_width
from mjlab.utils.spec import auto_wrap_fixed_base_mocap
from mjlab.utils.string import resolve_expr
from mjlab.utils.xml import fix_spec_xml, strip_buffer_textures

if TYPE_CHECKING:
  from mjlab.entity.variants import VariantMetadata


@dataclass(frozen=False)
class EntityIndexing:
  """Maps entity elements to global indices and addresses in the simulation."""

  # Elements.
  bodies: tuple[mujoco.MjsBody, ...]
  joints: tuple[mujoco.MjsJoint, ...]
  geoms: tuple[mujoco.MjsGeom, ...]
  sites: tuple[mujoco.MjsSite, ...]
  tendons: tuple[mujoco.MjsTendon, ...]
  cameras: tuple[mujoco.MjsCamera, ...]
  lights: tuple[mujoco.MjsLight, ...]
  materials: tuple[mujoco.MjsMaterial, ...]
  textures: tuple[mujoco.MjsTexture, ...]
  pairs: tuple[mujoco.MjsPair, ...]
  actuators: tuple[mujoco.MjsActuator, ...] | None

  # Indices.
  body_ids: torch.Tensor
  geom_ids: torch.Tensor
  site_ids: torch.Tensor
  tendon_ids: torch.Tensor
  cam_ids: torch.Tensor
  light_ids: torch.Tensor
  mat_ids: torch.Tensor
  tex_ids: torch.Tensor
  pair_ids: torch.Tensor
  ctrl_ids: torch.Tensor
  joint_ids: torch.Tensor
  mocap_id: int | None

  # Addresses.
  joint_q_adr: torch.Tensor
  joint_v_adr: torch.Tensor
  free_joint_q_adr: torch.Tensor
  free_joint_v_adr: torch.Tensor

  @property
  def root_body_id(self) -> int:
    return self.bodies[0].id


def _outer_index(
  env_ids: torch.Tensor | slice, ids: torch.Tensor | slice
) -> tuple[torch.Tensor | slice, torch.Tensor | slice]:
  """Make [env_ids, ids] select the outer product when both are tensors.

  Plain [tensor, tensor] indexing pairs the ids elementwise, which silently
  writes a diagonal whenever the shapes happen to broadcast.
  """
  if isinstance(env_ids, torch.Tensor) and isinstance(ids, torch.Tensor):
    return env_ids[:, None], ids
  return env_ids, ids


@dataclass
class EntityCfg:
  @dataclass
  class InitialStateCfg:
    # Root position and orientation.
    pos: tuple[float, float, float] = (0.0, 0.0, 0.0)
    rot: tuple[float, float, float, float] = (1.0, 0.0, 0.0, 0.0)
    # Root linear and angular velocity (only for floating base entities).
    lin_vel: tuple[float, float, float] = (0.0, 0.0, 0.0)
    ang_vel: tuple[float, float, float] = (0.0, 0.0, 0.0)
    # Articulation (only for articulated entities).
    # Set to None to use the model's existing keyframe (errors if none exists).
    joint_pos: dict[str, float] | None = field(default_factory=lambda: {".*": 0.0})
    joint_vel: dict[str, float] = field(default_factory=lambda: {".*": 0.0})

  init_state: InitialStateCfg = field(default_factory=InitialStateCfg)
  spec_fn: Callable[[], mujoco.MjSpec] = field(
    default_factory=lambda: lambda: mujoco.MjSpec()
  )
  articulation: EntityArticulationInfoCfg | None = None
  sort_actuators: bool = False
  """When True, reorder actuators so that ``model.ctrl`` follows joint/tendon/site
  definition order rather than the order actuators appear in the config. XML actuators
  are excluded from sorting and always retain their declaration order.
  """

  # Editors.
  lights: tuple[spec_cfg.LightCfg, ...] = field(default_factory=tuple)
  cameras: tuple[spec_cfg.CameraCfg, ...] = field(default_factory=tuple)
  textures: tuple[spec_cfg.TextureCfg, ...] = field(default_factory=tuple)
  materials: tuple[spec_cfg.MaterialCfg, ...] = field(default_factory=tuple)
  meshes: tuple[spec_cfg.MeshCfg, ...] = field(default_factory=tuple)
  geoms: tuple[spec_cfg.GeomCfg, ...] = field(default_factory=tuple)
  collisions: tuple[spec_cfg.CollisionCfg, ...] = field(default_factory=tuple)

  def build(self) -> Entity:
    """Build entity instance from this config.

    Override in subclasses to return custom Entity types.
    """
    return Entity(self)


@dataclass
class EntityArticulationInfoCfg:
  actuators: tuple[actuator.ActuatorCfg, ...] = field(default_factory=tuple)
  soft_joint_pos_limit_factor: float = 1.0


class Entity:
  """An entity represents a physical object in the simulation.

  Entity Type Matrix
  ==================
  MuJoCo entities can be categorized along two dimensions:

  1. Base Type:
    - Fixed Base: Entity is welded to the world (no freejoint)
    - Floating Base: Entity has 6 DOF movement (has freejoint)

  2. Articulation:
    - Non-articulated: No joints other than freejoint
    - Articulated: Has joints in kinematic tree (may or may not be actuated)

  Fixed non-articulated entities can optionally be mocap bodies, whereby their
  position and orientation can be set directly each timestep rather than being
  determined by physics. This property can be useful for creating props with
  adjustable position and orientation.

  Supported Combinations:
  ----------------------
  | Type                      | Example             | is_fixed_base | is_articulated | is_actuated |
  |---------------------------|---------------------|---------------|----------------|-------------|
  | Fixed Non-articulated     | Table, wall         | True          | False          | False       |
  | Fixed Articulated         | Robot arm, door     | True          | True           | True/False  |
  | Floating Non-articulated  | Box, ball, mug      | False         | False          | False       |
  | Floating Articulated      | Humanoid, quadruped | False         | True           | True/False  |
  """

  def __init__(self, cfg: EntityCfg) -> None:
    self.cfg = cfg
    self._actuators: list[actuator.Actuator] = []
    self._variant_metadata: VariantMetadata | None = None
    self._build_spec()
    self._identify_joints()
    self._apply_spec_editors()
    self._add_actuators()
    self._add_initial_state_keyframe()

  def _build_spec(self) -> None:
    from mjlab.entity.variants import VariantEntityCfg, build_merged_variant_spec

    if isinstance(self.cfg, VariantEntityCfg):
      self._spec, self._variant_metadata = build_merged_variant_spec(self.cfg)
    else:
      self._spec = auto_wrap_fixed_base_mocap(self.cfg.spec_fn)()

  @property
  def variant_metadata(self) -> VariantMetadata | None:
    return self._variant_metadata

  def _identify_joints(self) -> None:
    self._all_joints = self._spec.joints
    self._free_joint = None
    self._non_free_joints = tuple(self._all_joints)

    free_joints = [j for j in self._all_joints if j.type == mujoco.mjtJoint.mjJNT_FREE]
    if len(free_joints) > 1:
      raise ValueError(
        f"Entity spec has {len(free_joints)} freejoints. An Entity models a "
        "single rigid- or articulated-body system with at most one freejoint, "
        "which serves as its root. Model each detached floating body as its own "
        "entry in SceneCfg.entities instead."
      )

    if self._all_joints and self._all_joints[0].type == mujoco.mjtJoint.mjJNT_FREE:
      self._free_joint = self._all_joints[0]
      if not self._free_joint.name:
        self._free_joint.name = "floating_base_joint"
      self._non_free_joints = tuple(self._all_joints[1:])

  def _apply_spec_editors(self) -> None:
    spec_cfg.warn_overlapping_geom_edits(
      self.cfg.geoms, self.cfg.collisions, self._spec
    )
    for cfg_list in [
      self.cfg.lights,
      self.cfg.cameras,
      self.cfg.textures,
      self.cfg.materials,
      self.cfg.meshes,
      self.cfg.geoms,
      self.cfg.collisions,
    ]:
      for cfg in cfg_list:
        cfg.edit_spec(self._spec)

  def _add_actuators(self) -> None:
    if self.cfg.articulation is None:
      return

    # Collect actuator instances and their targets.
    pending: list[tuple[actuator.ActuatorCfg, actuator.Actuator, list[str]]] = []
    for actuator_cfg in self.cfg.articulation.actuators:
      # Find targets based on transmission type. resolve_matching_names raises
      # ValueError when no regex matches; we catch that to produce a better error with
      # namespace hints below.
      target_ids: list[int] = []
      target_names: list[str] = []
      target_spec_names: list[str] = []
      try:
        if actuator_cfg.transmission_type == TransmissionType.JOINT:
          target_ids, target_names = self.find_joints(actuator_cfg.target_names_expr)
          target_spec_names = [self._non_free_joints[i].name for i in target_ids]
        elif actuator_cfg.transmission_type == TransmissionType.TENDON:
          target_ids, target_names = self.find_tendons(actuator_cfg.target_names_expr)
          target_spec_names = [self._spec.tendons[i].name for i in target_ids]
        elif actuator_cfg.transmission_type == TransmissionType.SITE:
          target_ids, target_names = self.find_sites(actuator_cfg.target_names_expr)
          target_spec_names = [self.spec.sites[i].name for i in target_ids]
        else:
          raise TypeError(
            f"Invalid transmission_type: {actuator_cfg.transmission_type}. "
            f"Must be TransmissionType.JOINT, TransmissionType.TENDON, "
            f"or TransmissionType.SITE."
          )
      except ValueError:
        pass  # target_names stays empty, fall through to hint logic

      # Check other namespaces for matches. If we found nothing, this produces a
      # helpful error. If we did find targets, it warns about unactuated matches in
      # other namespaces.
      current = actuator_cfg.transmission_type
      other_matches: dict[TransmissionType, tuple[str, list[str]]] = {}
      other_namespaces = {
        TransmissionType.JOINT: ("joint", self.joint_names),
        TransmissionType.TENDON: ("tendon", self.tendon_names),
        TransmissionType.SITE: ("site", self.site_names),
      }
      for tt, (label, names) in other_namespaces.items():
        if tt == current or not names:
          continue
        try:
          _, matched = resolve_matching_names(actuator_cfg.target_names_expr, names)
          other_matches[tt] = (label, matched)
        except ValueError:
          pass

      if len(target_names) == 0:
        msg = (
          f"No {current.value}s matched expressions: {actuator_cfg.target_names_expr}"
        )
        if other_matches:
          hints = [
            f"{label}s ({', '.join(matched)})"
            for label, matched in other_matches.values()
          ]
          msg += (
            f". Matches were found in: {'; '.join(hints)}. "
            f"Check that transmission_type is correct."
          )
        raise ValueError(msg)

      for tt, (label, matched) in other_matches.items():
        warnings.warn(
          f"Actuator config matched {len(target_names)} {current.value}(s) "
          f"but the same expressions also match {len(matched)} {label}(s): "
          f"{', '.join(matched)}. Add a separate config with "
          f"transmission_type=TransmissionType.{tt.name} if those should "
          f"be actuated too.",
          stacklevel=2,
        )

      actuator_instance = actuator_cfg.build(self, target_ids, target_names)
      self._actuators.append(actuator_instance)
      pending.append((actuator_cfg, actuator_instance, target_spec_names))

    if not self.cfg.sort_actuators:
      for _, inst, names in pending:
        inst.edit_spec(self._spec, names)
      return

    # Sort actuators so ctrl order matches joint/tendon/site definition order.
    # XmlActuators are added first (they wrap pre-existing XML actuators),
    # then remaining actuators sorted by transmission type and target order.
    order_maps = {
      TransmissionType.JOINT: {name: i for i, name in enumerate(self.joint_names)},
      TransmissionType.TENDON: {name: i for i, name in enumerate(self.tendon_names)},
      TransmissionType.SITE: {name: i for i, name in enumerate(self.site_names)},
    }
    # Group by transmission type (ordering is conventional, not physics-motivated).
    # Within each group, actuators are sorted by their target's definition order in the
    # spec.
    type_priority = {
      TransmissionType.JOINT: 0,
      TransmissionType.TENDON: 1,
      TransmissionType.SITE: 2,
    }

    # XmlActuators go first in declaration order (they reference actuators already
    # present in the spec).
    for _, inst, names in pending:
      if isinstance(inst, XmlActuator):
        inst.edit_spec(self._spec, names)

    # Flatten remaining actuators to (instance, single_target) pairs and sort.
    flat: list[tuple[actuator.ActuatorCfg, actuator.Actuator, str]] = []
    for cfg, inst, names in pending:
      if not isinstance(inst, XmlActuator):
        for name in names:
          flat.append((cfg, inst, name))

    flat.sort(
      key=lambda item: (
        type_priority[item[0].transmission_type],
        order_maps[item[0].transmission_type].get(item[2], float("inf")),
      )
    )
    for _, inst, name in flat:
      inst.edit_spec(self._spec, [name])

  def _add_initial_state_keyframe(self) -> None:
    # If joint_pos is None, use existing keyframe from the model.
    if self.cfg.init_state.joint_pos is None:
      if not self._spec.keys:
        raise ValueError(
          "joint_pos=None requires the model to have a keyframe, but none exists."
        )
      # Keep the existing keyframe, just rename it.
      self._spec.keys[0].name = "init_state"
      if self.is_fixed_base:
        self.root_body.pos[:] = self.cfg.init_state.pos
        self.root_body.quat[:] = self.cfg.init_state.rot
      return

    qpos_components = []

    if self._free_joint is not None:
      qpos_components.extend([self.cfg.init_state.pos, self.cfg.init_state.rot])

    joint_pos = None
    if self._non_free_joints:
      joint_pos = resolve_expr(self.cfg.init_state.joint_pos, self.joint_names, 0.0)
      qpos_components.append(joint_pos)

    key_qpos = np.hstack(qpos_components) if qpos_components else np.array([])
    key = self._spec.add_key(name="init_state", qpos=key_qpos.tolist())

    if self.is_actuated and joint_pos is not None:
      name_to_pos = {name: joint_pos[i] for i, name in enumerate(self.joint_names)}
      ctrl = []
      for act in self._spec.actuators:
        joint_name = act.target
        ctrl.append(name_to_pos.get(joint_name, 0.0))
      key.ctrl = np.array(ctrl)

    if self.is_fixed_base:
      self.root_body.pos[:] = self.cfg.init_state.pos
      self.root_body.quat[:] = self.cfg.init_state.rot

  # Attributes.

  @property
  def is_fixed_base(self) -> bool:
    """Entity is welded to the world."""
    return self._free_joint is None

  @property
  def is_articulated(self) -> bool:
    """Entity is articulated (has fixed or actuated joints)."""
    return len(self._non_free_joints) > 0

  @property
  def is_actuated(self) -> bool:
    """Entity has actuated joints."""
    return len(self._actuators) > 0

  @property
  def has_tendon_actuators(self) -> bool:
    """Entity has actuators using tendon transmission."""
    if self.cfg.articulation is None:
      return False
    return any(
      act.transmission_type == TransmissionType.TENDON
      for act in self.cfg.articulation.actuators
    )

  @property
  def has_site_actuators(self) -> bool:
    """Entity has actuators using site transmission."""
    if self.cfg.articulation is None:
      return False
    return any(
      act.transmission_type == TransmissionType.SITE
      for act in self.cfg.articulation.actuators
    )

  @property
  def is_mocap(self) -> bool:
    """Entity root body is a mocap body (only for fixed-base entities)."""
    return bool(self.root_body.mocap) if self.is_fixed_base else False

  @property
  def spec(self) -> mujoco.MjSpec:
    return self._spec

  @property
  def data(self) -> EntityData:
    return self._data

  @property
  def actuators(self) -> list[actuator.Actuator]:
    return self._actuators

  # Names.

  @property
  def body_names(self) -> tuple[str, ...]:
    return tuple(b.name.split("/")[-1] for b in self.spec.bodies[1:])

  @property
  def all_joint_names(self) -> tuple[str, ...]:
    return tuple(j.name.split("/")[-1] for j in self._all_joints)

  @property
  def joint_names(self) -> tuple[str, ...]:
    return tuple(j.name.split("/")[-1] for j in self._non_free_joints)

  @property
  def geom_names(self) -> tuple[str, ...]:
    return tuple(g.name.split("/")[-1] for g in self.spec.geoms)

  @property
  def site_names(self) -> tuple[str, ...]:
    return tuple(s.name.split("/")[-1] for s in self.spec.sites)

  @property
  def tendon_names(self) -> tuple[str, ...]:
    return tuple(t.name.split("/")[-1] for t in self._spec.tendons)

  @property
  def camera_names(self) -> tuple[str, ...]:
    return tuple(c.name.split("/")[-1] for c in self.spec.cameras)

  @property
  def light_names(self) -> tuple[str, ...]:
    return tuple(lt.name.split("/")[-1] for lt in self.spec.lights)

  @property
  def material_names(self) -> tuple[str, ...]:
    return tuple(m.name.split("/")[-1] for m in self.spec.materials)

  @property
  def texture_names(self) -> tuple[str, ...]:
    return tuple(t.name.split("/")[-1] for t in self.spec.textures)

  @property
  def pair_names(self) -> tuple[str, ...]:
    return tuple(p.name.split("/")[-1] for p in self.spec.pairs)

  @property
  def actuator_names(self) -> tuple[str, ...]:
    return tuple(a.name.split("/")[-1] for a in self.spec.actuators)

  # Counts.

  @property
  def num_bodies(self) -> int:
    return len(self.body_names)

  @property
  def num_joints(self) -> int:
    return len(self.joint_names)

  @property
  def num_geoms(self) -> int:
    return len(self.geom_names)

  @property
  def num_sites(self) -> int:
    return len(self.site_names)

  @property
  def num_tendons(self) -> int:
    return len(self.tendon_names)

  @property
  def num_cameras(self) -> int:
    return len(self.camera_names)

  @property
  def num_lights(self) -> int:
    return len(self.light_names)

  @property
  def num_materials(self) -> int:
    return len(self.material_names)

  @property
  def num_textures(self) -> int:
    return len(self.texture_names)

  @property
  def num_pairs(self) -> int:
    return len(self.pair_names)

  @property
  def num_actuators(self) -> int:
    return len(self.actuator_names)

  @property
  def root_body(self) -> mujoco.MjsBody:
    return self.spec.bodies[1]

  # Find methods.

  def find_bodies(
    self, name_keys: str | Sequence[str], preserve_order: bool = False
  ) -> tuple[list[int], list[str]]:
    return resolve_matching_names(name_keys, self.body_names, preserve_order)

  def find_joints(
    self,
    name_keys: str | Sequence[str],
    joint_subset: Sequence[str] | None = None,
    preserve_order: bool = False,
  ) -> tuple[list[int], list[str]]:
    if joint_subset is None:
      joint_subset = self.joint_names
    return resolve_matching_names(name_keys, joint_subset, preserve_order)

  def find_joints_by_actuator_names(
    self,
    actuator_name_keys: str | Sequence[str],
  ) -> tuple[list[int], list[str]]:
    # Collect all actuated joint names.
    actuated_joint_names_set = set()
    for act in self._actuators:
      actuated_joint_names_set.update(act.target_names)

    # Filter self.joint_names to only actuated joints, preserving natural order.
    actuated_in_natural_order = [
      name for name in self.joint_names if name in actuated_joint_names_set
    ]

    # Find joints matching the pattern within actuated joints.
    _, matched_joint_names = self.find_joints(
      actuator_name_keys, joint_subset=actuated_in_natural_order, preserve_order=False
    )

    # Map joint names back to entity-local indices (indices into self.joint_names).
    name_to_entity_idx = {name: i for i, name in enumerate(self.joint_names)}
    joint_ids = [name_to_entity_idx[name] for name in matched_joint_names]
    return joint_ids, matched_joint_names

  def find_geoms(
    self,
    name_keys: str | Sequence[str],
    geom_subset: Sequence[str] | None = None,
    preserve_order: bool = False,
  ) -> tuple[list[int], list[str]]:
    if geom_subset is None:
      geom_subset = self.geom_names
    return resolve_matching_names(name_keys, geom_subset, preserve_order)

  def find_sites(
    self,
    name_keys: str | Sequence[str],
    site_subset: Sequence[str] | None = None,
    preserve_order: bool = False,
  ) -> tuple[list[int], list[str]]:
    if site_subset is None:
      site_subset = self.site_names
    return resolve_matching_names(name_keys, site_subset, preserve_order)

  def find_tendons(
    self,
    name_keys: str | Sequence[str],
    tendon_subset: Sequence[str] | None = None,
    preserve_order: bool = False,
  ) -> tuple[list[int], list[str]]:
    if tendon_subset is None:
      tendon_subset = self.tendon_names
    return resolve_matching_names(name_keys, tendon_subset, preserve_order)

  def find_cameras(
    self,
    name_keys: str | Sequence[str],
    camera_subset: Sequence[str] | None = None,
    preserve_order: bool = False,
  ) -> tuple[list[int], list[str]]:
    if camera_subset is None:
      camera_subset = self.camera_names
    return resolve_matching_names(name_keys, camera_subset, preserve_order)

  def find_lights(
    self,
    name_keys: str | Sequence[str],
    light_subset: Sequence[str] | None = None,
    preserve_order: bool = False,
  ) -> tuple[list[int], list[str]]:
    if light_subset is None:
      light_subset = self.light_names
    return resolve_matching_names(name_keys, light_subset, preserve_order)

  def find_materials(
    self,
    name_keys: str | Sequence[str],
    material_subset: Sequence[str] | None = None,
    preserve_order: bool = False,
  ) -> tuple[list[int], list[str]]:
    if material_subset is None:
      material_subset = self.material_names
    return resolve_matching_names(name_keys, material_subset, preserve_order)

  def find_textures(
    self,
    name_keys: str | Sequence[str],
    texture_subset: Sequence[str] | None = None,
    preserve_order: bool = False,
  ) -> tuple[list[int], list[str]]:
    if texture_subset is None:
      texture_subset = self.texture_names
    return resolve_matching_names(name_keys, texture_subset, preserve_order)

  def find_pairs(
    self,
    name_keys: str | Sequence[str],
    pair_subset: Sequence[str] | None = None,
    preserve_order: bool = False,
  ) -> tuple[list[int], list[str]]:
    if pair_subset is None:
      pair_subset = self.pair_names
    return resolve_matching_names(name_keys, pair_subset, preserve_order)

  def find_actuators(
    self,
    name_keys: str | Sequence[str],
    actuator_subset: Sequence[str] | None = None,
    preserve_order: bool = False,
  ) -> tuple[list[int], list[str]]:
    if actuator_subset is None:
      actuator_subset = self.actuator_names
    return resolve_matching_names(name_keys, actuator_subset, preserve_order)

  def compile(self) -> mujoco.MjModel:
    """Compile the underlying MjSpec into an MjModel."""
    return self.spec.compile()

  def write_xml(self, xml_path: Path) -> None:
    """Write the MjSpec to disk.

    Operates on a copy of the spec to avoid mutating the original.
    """
    tmp = self.spec.copy()
    strip_buffer_textures(tmp)
    xml_path.write_text(fix_spec_xml(tmp.to_xml()))

  def to_zip(self, path: Path) -> None:
    """Write the MjSpec to a zip file."""
    with path.open("wb") as f:
      mujoco.MjSpec.to_zip(self.spec, f)

  def initialize(
    self,
    mj_model: mujoco.MjModel,
    model: mjwarp.Model,
    data: mjwarp.Data,
    device: str,
  ) -> None:
    """Prepare the entity for simulation after the spec has been compiled.

    Computes global index mappings, initializes actuators, and allocates all nworld
    state and target tensors in ``EntityData``. Called once by the scene during
    environment construction.
    """
    indexing = self._compute_indexing(mj_model, device)
    self.indexing = indexing
    nworld = data.nworld

    for act in self._actuators:
      act.initialize(mj_model, model, data, device)

    # Vectorize built-in actuators, then fuse ideal PD actuators; we'll loop
    # through whatever custom actuators remain.
    builtin_group, custom_actuators = BuiltinActuatorGroup.process(self._actuators)
    builtin_group.initialize(nworld, device)
    self._builtin_group = builtin_group
    fused_actuator_group, custom_actuators = FusedActuatorGroup.process(
      custom_actuators
    )
    fused_actuator_group.initialize(nworld, device)
    self._fused_actuator_group = fused_actuator_group
    self._custom_actuators = custom_actuators

    # Root state.
    default_root_state = self._build_default_root_state(nworld, device)

    # Joint state.
    if self.is_articulated:
      if self.cfg.init_state.joint_pos is None:
        # Use keyframe joint positions.
        key_qpos = mj_model.key("init_state").qpos
        nq_root = 7 if not self.is_fixed_base else 0
        default_joint_pos = torch.tensor(key_qpos[nq_root:], device=device)[
          None
        ].repeat(nworld, 1)
      else:
        default_joint_pos = torch.tensor(
          resolve_expr(self.cfg.init_state.joint_pos, self.joint_names, 0.0),
          device=device,
        )[None].repeat(nworld, 1)
      default_joint_vel = torch.tensor(
        resolve_expr(self.cfg.init_state.joint_vel, self.joint_names, 0.0),
        device=device,
      )[None].repeat(nworld, 1)

      # Joint limits.
      joint_ids_list = [j.id for j in self._non_free_joints]
      dof_limits = model.jnt_range[:, joint_ids_list]
      default_joint_pos_limits = dof_limits.clone()
      joint_pos_limits = default_joint_pos_limits.clone()

      joint_pos_mean = (joint_pos_limits[..., 0] + joint_pos_limits[..., 1]) / 2
      joint_pos_range = joint_pos_limits[..., 1] - joint_pos_limits[..., 0]

      # Soft limits.
      soft_limit_factor = (
        self.cfg.articulation.soft_joint_pos_limit_factor
        if self.cfg.articulation
        else 1.0
      )
      soft_joint_pos_limits = torch.stack(
        [
          joint_pos_mean - 0.5 * joint_pos_range * soft_limit_factor,
          joint_pos_mean + 0.5 * joint_pos_range * soft_limit_factor,
        ],
        dim=-1,
      )

      # Unlimited joints have jnt_range=[0,0] in MuJoCo, which makes all
      # the computed limits [0,0]. Override to [-inf, inf] so downstream
      # clamping becomes a no-op. (Can't do this before soft-limit math
      # because inf - inf = NaN.)
      unlimited = ~torch.tensor(
        mj_model.jnt_limited[joint_ids_list], device=device, dtype=torch.bool
      )
      for limits in (joint_pos_limits, default_joint_pos_limits, soft_joint_pos_limits):
        limits[:, unlimited, 0] = float("-inf")
        limits[:, unlimited, 1] = float("inf")
    else:
      empty_shape = (nworld, 0)
      default_joint_pos = torch.empty(*empty_shape, dtype=torch.float, device=device)
      default_joint_vel = torch.empty(*empty_shape, dtype=torch.float, device=device)
      default_joint_pos_limits = torch.empty(
        *empty_shape, 2, dtype=torch.float, device=device
      )
      joint_pos_limits = torch.empty(*empty_shape, 2, dtype=torch.float, device=device)
      soft_joint_pos_limits = torch.empty(
        *empty_shape, 2, dtype=torch.float, device=device
      )

    if self.is_actuated:
      joint_pos_target = torch.zeros(
        (nworld, self.num_joints), dtype=torch.float, device=device
      )
      joint_vel_target = torch.zeros(
        (nworld, self.num_joints), dtype=torch.float, device=device
      )
      joint_effort_target = torch.zeros(
        (nworld, self.num_joints), dtype=torch.float, device=device
      )
    else:
      joint_pos_target = torch.empty(nworld, 0, dtype=torch.float, device=device)
      joint_vel_target = torch.empty(nworld, 0, dtype=torch.float, device=device)
      joint_effort_target = torch.empty(nworld, 0, dtype=torch.float, device=device)

    # Only allocate tendon targets if there are actuators using tendon transmission.
    if self.has_tendon_actuators:
      num_tendons = len(self.tendon_names)
      tendon_len_target = torch.zeros(
        (nworld, num_tendons), dtype=torch.float, device=device
      )
      tendon_vel_target = torch.zeros(
        (nworld, num_tendons), dtype=torch.float, device=device
      )
      tendon_effort_target = torch.zeros(
        (nworld, num_tendons), dtype=torch.float, device=device
      )
    else:
      tendon_len_target = torch.empty(nworld, 0, dtype=torch.float, device=device)
      tendon_vel_target = torch.empty(nworld, 0, dtype=torch.float, device=device)
      tendon_effort_target = torch.empty(nworld, 0, dtype=torch.float, device=device)

    # Only allocate site targets if there are actuators using site transmission.
    if self.has_site_actuators:
      num_sites = len(self.site_names)
      site_effort_target = torch.zeros(
        (nworld, num_sites), dtype=torch.float, device=device
      )
    else:
      site_effort_target = torch.empty(nworld, 0, dtype=torch.float, device=device)

    # Encoder bias for simulating encoder calibration errors.
    # Shape: (num_envs, num_joints). Defaults to zero (no bias).
    if self.is_articulated:
      encoder_bias = torch.zeros(
        (nworld, self.num_joints), dtype=torch.float, device=device
      )
    else:
      encoder_bias = torch.empty(nworld, 0, dtype=torch.float, device=device)

    self._data = EntityData(
      indexing=indexing,
      data=data,
      model=model,
      device=device,
      default_root_state=default_root_state,
      default_joint_pos=default_joint_pos,
      default_joint_vel=default_joint_vel,
      default_joint_pos_limits=default_joint_pos_limits,
      joint_pos_limits=joint_pos_limits,
      soft_joint_pos_limits=soft_joint_pos_limits,
      gravity_vec_w=torch.tensor([0.0, 0.0, -1.0], device=device).repeat(nworld, 1),
      forward_vec_b=torch.tensor([1.0, 0.0, 0.0], device=device).repeat(nworld, 1),
      is_fixed_base=self.is_fixed_base,
      is_articulated=self.is_articulated,
      is_actuated=self.is_actuated,
      joint_pos_target=joint_pos_target,
      joint_vel_target=joint_vel_target,
      joint_effort_target=joint_effort_target,
      tendon_len_target=tendon_len_target,
      tendon_vel_target=tendon_vel_target,
      tendon_effort_target=tendon_effort_target,
      site_effort_target=site_effort_target,
      encoder_bias=encoder_bias,
    )

  def update(self, dt: float) -> None:
    """Advance actuator internal state by one physics substep.

    Called after each ``sim.step()`` within the decimation loop.
    """
    for act in self._actuators:
      act.update(dt)

  def reset(self, env_ids: torch.Tensor | slice | None = None) -> None:
    """Zero actuator targets and reset actuator internal state.

    Called by the scene when environments are reset at episode boundaries,
    and by commands that teleport the robot to a new pose mid-episode.
    """
    self._data.clear_state(env_ids)

    for act in self._actuators:
      act.reset(env_ids)

  def clear_state(self, env_ids: torch.Tensor | slice | None = None) -> None:
    """Deprecated. Use ``reset`` instead."""
    warnings.warn(
      "Entity.clear_state() is deprecated. Use Entity.reset().",
      DeprecationWarning,
      stacklevel=2,
    )
    self.reset(env_ids)

  def write_data_to_sim(self) -> None:
    """Convert actuator targets into low-level controls and write them to the sim.

    Called before each ``sim.step()`` within the decimation loop. Builtin actuators are
    applied in a single batched operation; custom actuators are applied individually.
    """
    self._apply_actuator_controls()

  def write_ctrl_to_sim(
    self,
    ctrl: torch.Tensor,
    ctrl_ids: torch.Tensor | slice | None = None,
    env_ids: torch.Tensor | slice | None = None,
  ) -> None:
    """Write control inputs to the simulation.

    Args:
      ctrl: A tensor of control inputs.
      ctrl_ids: A tensor of control indices.
      env_ids: Optional tensor or slice specifying which environments to set.
        If None, all environments are set.
    """
    self._data.write_ctrl(ctrl, ctrl_ids, env_ids)

  def write_root_state_to_sim(
    self, root_state: torch.Tensor, env_ids: torch.Tensor | slice | None = None
  ) -> None:
    """Set the root state into the simulation.

    The root state consists of position (3), orientation as a (w, x, y, z)
    quaternion (4), linear velocity (3), and angular velocity (3), for a total
    of 13 values. All of the quantities are in the world frame.

    Args:
      root_state: Tensor of shape (N, 13) where N is the number of environments.
      env_ids: Optional tensor or slice specifying which environments to set. If
        None, all environments are set.
    """
    self._data.write_root_state(root_state, env_ids)

  def write_root_link_pose_to_sim(
    self,
    root_pose: torch.Tensor,
    env_ids: torch.Tensor | slice | None = None,
  ):
    """Set the root pose into the simulation. Like `write_root_state_to_sim()`
    but only sets position and orientation.

    Args:
      root_pose: Tensor of shape (N, 7) where N is the number of environments.
      env_ids: Optional tensor or slice specifying which environments to set. If
        None, all environments are set.
    """
    self._data.write_root_pose(root_pose, env_ids)

  def write_root_link_velocity_to_sim(
    self,
    root_velocity: torch.Tensor,
    env_ids: torch.Tensor | slice | None = None,
  ):
    """Set the root link (body origin) velocity into the simulation. Like
    `write_root_state_to_sim()` but only sets linear and angular velocity.

    Args:
      root_velocity: Tensor of shape (N, 6) where N is the number of environments.
        Contains linear velocity (3) at body origin and angular velocity (3),
        both in world frame.
      env_ids: Optional tensor or slice specifying which environments to set. If
        None, all environments are set.
    """
    self._data.write_root_velocity(root_velocity, env_ids)

  def write_root_link_velocity_b_to_sim(
    self,
    root_velocity_b: torch.Tensor,
    env_ids: torch.Tensor | slice | None = None,
  ):
    """Like `write_root_link_velocity_to_sim()` but the velocity is expressed
    in the root link's body frame. Reads the orientation from qpos, so it is
    safe to call during a reset before forward() runs.

    Args:
      root_velocity_b: Tensor of shape (N, 6) where N is the number of
        environments. Contains linear velocity (3) at body origin and angular
        velocity (3), both in the root link's body frame.
      env_ids: Optional tensor or slice specifying which environments to set. If
        None, all environments are set.
    """
    self._data.write_root_velocity_b(root_velocity_b, env_ids)

  def write_root_com_velocity_to_sim(
    self,
    root_velocity: torch.Tensor,
    env_ids: torch.Tensor | slice | None = None,
  ):
    """Set the root COM velocity into the simulation.

    Args:
      root_velocity: Tensor of shape (N, 6) where N is the number of environments.
        Contains linear velocity (3) at COM and angular velocity (3),
        both in world frame.
      env_ids: Optional tensor or slice specifying which environments to set. If
        None, all environments are set.
    """
    self._data.write_root_com_velocity(root_velocity, env_ids)

  def write_joint_state_to_sim(
    self,
    position: torch.Tensor,
    velocity: torch.Tensor,
    joint_ids: torch.Tensor | slice | None = None,
    env_ids: torch.Tensor | slice | None = None,
  ):
    """Set the joint state into the simulation.

    The joint state consists of joint positions and velocities. It does not include
    the root state.

    Args:
      position: Tensor of shape (N, num_joints) where N is the number of environments.
      velocity: Tensor of shape (N, num_joints) where N is the number of environments.
      joint_ids: Optional tensor or slice specifying which joints to set. If None,
        all joints are set.
      env_ids: Optional tensor or slice specifying which environments to set. If
        None, all environments are set.
    """
    self._data.write_joint_state(position, velocity, joint_ids, env_ids)

  def write_joint_position_to_sim(
    self,
    position: torch.Tensor,
    joint_ids: torch.Tensor | slice | None = None,
    env_ids: torch.Tensor | slice | None = None,
  ):
    """Set the joint positions into the simulation. Like `write_joint_state_to_sim()`
    but only sets joint positions.

    Args:
      position: Tensor of shape (N, num_joints) where N is the number of environments.
      joint_ids: Optional tensor or slice specifying which joints to set. If None,
        all joints are set.
      env_ids: Optional tensor or slice specifying which environments to set. If
        None, all environments are set.
    """
    self._data.write_joint_position(position, joint_ids, env_ids)

  def write_joint_velocity_to_sim(
    self,
    velocity: torch.Tensor,
    joint_ids: torch.Tensor | slice | None = None,
    env_ids: torch.Tensor | slice | None = None,
  ):
    """Set the joint velocities into the simulation. Like `write_joint_state_to_sim()`
    but only sets joint velocities.

    Args:
      velocity: Tensor of shape (N, num_joints) where N is the number of environments.
      joint_ids: Optional tensor or slice specifying which joints to set. If None,
        all joints are set.
      env_ids: Optional tensor or slice specifying which environments to set. If
        None, all environments are set.
    """
    self._data.write_joint_velocity(velocity, joint_ids, env_ids)

  def set_joint_position_target(
    self,
    position: torch.Tensor,
    joint_ids: torch.Tensor | slice | None = None,
    env_ids: torch.Tensor | slice | None = None,
  ) -> None:
    """Set joint position targets.

    Args:
      position: Target joint poisitions with shape (N, num_joints).
      joint_ids: Optional joint indices to set. If None, set all joints.
      env_ids: Optional environment indices. If None, set all environments.
    """
    if env_ids is None:
      env_ids = slice(None)
    if joint_ids is None:
      joint_ids = slice(None)
    env_ids, joint_ids = _outer_index(env_ids, joint_ids)
    self._data.joint_pos_target[env_ids, joint_ids] = position

  def set_joint_velocity_target(
    self,
    velocity: torch.Tensor,
    joint_ids: torch.Tensor | slice | None = None,
    env_ids: torch.Tensor | slice | None = None,
  ) -> None:
    """Set joint velocity targets.

    Args:
      velocity: Target joint velocities with shape (N, num_joints).
      joint_ids: Optional joint indices to set. If None, set all joints.
      env_ids: Optional environment indices. If None, set all environments.
    """
    if env_ids is None:
      env_ids = slice(None)
    if joint_ids is None:
      joint_ids = slice(None)
    env_ids, joint_ids = _outer_index(env_ids, joint_ids)
    self._data.joint_vel_target[env_ids, joint_ids] = velocity

  def set_joint_effort_target(
    self,
    effort: torch.Tensor,
    joint_ids: torch.Tensor | slice | None = None,
    env_ids: torch.Tensor | slice | None = None,
  ) -> None:
    """Set joint effort targets.

    Args:
      effort: Target joint efforts with shape (N, num_joints).
      joint_ids: Optional joint indices to set. If None, set all joints.
      env_ids: Optional environment indices. If None, set all environments.
    """
    if env_ids is None:
      env_ids = slice(None)
    if joint_ids is None:
      joint_ids = slice(None)
    env_ids, joint_ids = _outer_index(env_ids, joint_ids)
    self._data.joint_effort_target[env_ids, joint_ids] = effort

  def set_tendon_len_target(
    self,
    length: torch.Tensor,
    tendon_ids: torch.Tensor | slice | None = None,
    env_ids: torch.Tensor | slice | None = None,
  ) -> None:
    """Set tendon length targets.

    Args:
      length: Target tendon lengths with shape (N, num_tendons).
      tendon_ids: Optional tendon indices to set. If None, set all tendons.
      env_ids: Optional environment indices. If None, set all environments.
    """
    if env_ids is None:
      env_ids = slice(None)
    if tendon_ids is None:
      tendon_ids = slice(None)
    env_ids, tendon_ids = _outer_index(env_ids, tendon_ids)
    self._data.tendon_len_target[env_ids, tendon_ids] = length

  def set_tendon_vel_target(
    self,
    velocity: torch.Tensor,
    tendon_ids: torch.Tensor | slice | None = None,
    env_ids: torch.Tensor | slice | None = None,
  ) -> None:
    """Set tendon velocity targets.

    Args:
      velocity: Target tendon velocities with shape (N, num_tendons).
      tendon_ids: Optional tendon indices to set. If None, set all tendons.
      env_ids: Optional environment indices. If None, set all environments.
    """
    if env_ids is None:
      env_ids = slice(None)
    if tendon_ids is None:
      tendon_ids = slice(None)
    env_ids, tendon_ids = _outer_index(env_ids, tendon_ids)
    self._data.tendon_vel_target[env_ids, tendon_ids] = velocity

  def set_tendon_effort_target(
    self,
    effort: torch.Tensor,
    tendon_ids: torch.Tensor | slice | None = None,
    env_ids: torch.Tensor | slice | None = None,
  ) -> None:
    """Set tendon effort targets.

    Args:
      effort: Target tendon efforts with shape (N, num_tendons).
      tendon_ids: Optional tendon indices to set. If None, set all tendons.
      env_ids: Optional environment indices. If None, set all environments.
    """
    if env_ids is None:
      env_ids = slice(None)
    if tendon_ids is None:
      tendon_ids = slice(None)
    env_ids, tendon_ids = _outer_index(env_ids, tendon_ids)
    self._data.tendon_effort_target[env_ids, tendon_ids] = effort

  def set_site_effort_target(
    self,
    effort: torch.Tensor,
    site_ids: torch.Tensor | slice | None = None,
    env_ids: torch.Tensor | slice | None = None,
  ) -> None:
    """Set site effort targets.

    Args:
      effort: Target site efforts with shape (N, num_sites).
      site_ids: Optional site indices to set. If None, set all sites.
      env_ids: Optional environment indices. If None, set all environments.
    """
    if env_ids is None:
      env_ids = slice(None)
    if site_ids is None:
      site_ids = slice(None)
    env_ids, site_ids = _outer_index(env_ids, site_ids)
    self._data.site_effort_target[env_ids, site_ids] = effort

  def write_external_wrench_to_sim(
    self,
    forces: torch.Tensor,
    torques: torch.Tensor,
    env_ids: torch.Tensor | slice | None = None,
    body_ids: Sequence[int] | slice | None = None,
  ) -> None:
    """Apply external wrenches to bodies in the simulation.

    Underneath the hood, this sets the `xfrc_applied` field in the MuJoCo data
    structure. The wrenches are specified in the world frame and persist until
    the next call to this function or until the simulation is reset.

    Args:
      forces: Tensor of shape (N, num_bodies, 3) where N is the number of
        environments.
      torques: Tensor of shape (N, num_bodies, 3) where N is the number of
        environments.
      env_ids: Optional tensor or slice specifying which environments to set. If
        None, all environments are set.
      body_ids: Optional list of body indices or slice specifying which bodies to
        apply the wrenches to. If None, wrenches are applied to all bodies.
    """
    self._data.write_external_wrench(forces, torques, body_ids, env_ids)

  def write_mocap_pose_to_sim(
    self,
    mocap_pose: torch.Tensor,
    env_ids: torch.Tensor | slice | None = None,
  ) -> None:
    """Set the pose of a mocap body into the simulation.

    Args:
      mocap_pose: Tensor of shape (N, 7) where N is the number of environments.
        Format: [pos_x, pos_y, pos_z, quat_w, quat_x, quat_y, quat_z]
      env_ids: Optional tensor or slice specifying which environments to set. If
        None, all environments are set.
    """
    self._data.write_mocap_pose(mocap_pose, env_ids)

  ##
  # Private methods.
  ##

  def _build_default_root_state(self, nworld: int, device: str) -> torch.Tensor:
    """Build default root state tensor, uniform across all worlds."""
    base = self.cfg.init_state
    components: list[tuple[float, ...]] = [base.pos, base.rot]
    if not self.is_fixed_base:
      components.extend([base.lin_vel, base.ang_vel])
    return torch.tensor(
      sum((tuple(c) for c in components), ()),
      dtype=torch.float,
      device=device,
    ).repeat(nworld, 1)

  def _compute_indexing(self, model: mujoco.MjModel, device: str) -> EntityIndexing:
    bodies = tuple([b for b in self.spec.bodies[1:]])
    joints = self._non_free_joints
    geoms = tuple(self.spec.geoms)
    sites = tuple(self.spec.sites)
    tendons = tuple(self.spec.tendons)
    cameras = tuple(self.spec.cameras)
    lights = tuple(self.spec.lights)
    materials = tuple(self.spec.materials)
    textures = tuple(self.spec.textures)
    pairs = tuple(self.spec.pairs)

    body_ids = torch.tensor([b.id for b in bodies], dtype=torch.int, device=device)
    geom_ids = torch.tensor([g.id for g in geoms], dtype=torch.int, device=device)
    site_ids = torch.tensor([s.id for s in sites], dtype=torch.int, device=device)
    tendon_ids = torch.tensor([t.id for t in tendons], dtype=torch.int, device=device)
    cam_ids = torch.tensor([c.id for c in cameras], dtype=torch.int, device=device)
    light_ids = torch.tensor([lt.id for lt in lights], dtype=torch.int, device=device)
    mat_ids = torch.tensor([m.id for m in materials], dtype=torch.int, device=device)
    tex_ids = torch.tensor([t.id for t in textures], dtype=torch.int, device=device)
    pair_ids = torch.tensor([p.id for p in pairs], dtype=torch.int, device=device)
    joint_ids = torch.tensor([j.id for j in joints], dtype=torch.int, device=device)

    if self.is_actuated:
      actuators = tuple(self.spec.actuators)
      ctrl_ids = torch.tensor([a.id for a in actuators], dtype=torch.int, device=device)
    else:
      actuators = None
      ctrl_ids = torch.empty(0, dtype=torch.int, device=device)

    joint_q_adr = []
    joint_v_adr = []
    free_joint_q_adr = []
    free_joint_v_adr = []
    for joint in self.spec.joints:
      jnt = model.joint(joint.name)
      jnt_type = jnt.type[0]
      vadr = jnt.dofadr[0]
      qadr = jnt.qposadr[0]
      if jnt_type == mujoco.mjtJoint.mjJNT_FREE:
        free_joint_v_adr.extend(range(vadr, vadr + 6))
        free_joint_q_adr.extend(range(qadr, qadr + 7))
      else:
        joint_v_adr.extend(range(vadr, vadr + dof_width(jnt_type)))
        joint_q_adr.extend(range(qadr, qadr + qpos_width(jnt_type)))
    joint_q_adr = torch.tensor(joint_q_adr, dtype=torch.int, device=device)
    joint_v_adr = torch.tensor(joint_v_adr, dtype=torch.int, device=device)
    free_joint_v_adr = torch.tensor(free_joint_v_adr, dtype=torch.int, device=device)
    free_joint_q_adr = torch.tensor(free_joint_q_adr, dtype=torch.int, device=device)

    if self.is_fixed_base and self.is_mocap:
      mocap_id = int(model.body_mocapid[self.root_body.id])
    else:
      mocap_id = None

    return EntityIndexing(
      bodies=bodies,
      joints=joints,
      geoms=geoms,
      sites=sites,
      tendons=tendons,
      cameras=cameras,
      lights=lights,
      materials=materials,
      textures=textures,
      pairs=pairs,
      actuators=actuators,
      body_ids=body_ids,
      geom_ids=geom_ids,
      site_ids=site_ids,
      tendon_ids=tendon_ids,
      cam_ids=cam_ids,
      light_ids=light_ids,
      mat_ids=mat_ids,
      tex_ids=tex_ids,
      pair_ids=pair_ids,
      ctrl_ids=ctrl_ids,
      joint_ids=joint_ids,
      mocap_id=mocap_id,
      joint_q_adr=joint_q_adr,
      joint_v_adr=joint_v_adr,
      free_joint_q_adr=free_joint_q_adr,
      free_joint_v_adr=free_joint_v_adr,
    )

  def _apply_actuator_controls(self) -> None:
    self._builtin_group.apply_controls(self._data)
    self._fused_actuator_group.apply_controls(self._data)
    for act in self._custom_actuators:
      command = act.get_command(self._data)
      command = act.apply_delay(command)
      self._data.write_ctrl(act.compute(command), act.ctrl_ids)
