"""Tests for the dr.* model-field randomization engine."""

import math

import mujoco
import pytest
import torch
from conftest import get_test_device

from mjlab.entity import EntityCfg
from mjlab.envs.mdp import dr
from mjlab.managers.scene_entity_config import SceneEntityCfg
from mjlab.scene import Scene, SceneCfg
from mjlab.sim.sim import Simulation, SimulationCfg
from mjlab.utils.lab_api.math import quat_from_euler_xyz, quat_mul

pytestmark = pytest.mark.filterwarnings(
  "ignore:Use of index_put_ on expanded tensors is deprecated:UserWarning"
)

# Shared helpers.

ROBOT_XML = """
<mujoco>
  <worldbody>
    <body name="base" pos="0 0 1">
      <freejoint name="free_joint"/>
      <geom name="base_geom" type="box" size="0.1 0.1 0.1" mass="1.0"
        friction="0.5 0.01 0.005"/>
      <site name="base_site" pos="0 0 0.1"/>
      <camera name="front_cam" pos="0.5 0 0.3" xyaxes="0 -1 0 0 0 1" fovy="60"/>
      <camera name="side_cam" pos="0 0.5 0.3" xyaxes="1 0 0 0 0 1" fovy="45"/>
      <light name="top_light" pos="0 0 2" dir="0 0 -1"/>
      <body name="foot1" pos="0.2 0 0">
        <joint name="joint1" type="hinge" axis="0 0 1" range="0 1.57"/>
        <geom name="foot1_geom" type="box" size="0.05 0.05 0.05" mass="0.1"
          friction="0.5 0.01 0.005"/>
      </body>
      <body name="foot2" pos="-0.2 0 0">
        <joint name="joint2" type="hinge" axis="0 0 1" range="0 1.57"/>
        <geom name="foot2_geom" type="box" size="0.05 0.05 0.05" mass="0.1"
          friction="0.5 0.01 0.005"/>
      </body>
      <body name="offset_body" pos="0 0.3 0">
        <joint name="joint3" type="hinge" axis="1 0 0" range="-1.57 1.57"/>
        <inertial pos="0.05 -0.02 0.03" mass="0.5"
          fullinertia="0.001 0.002 0.0015 0.0001 -0.00005 0.00008"/>
        <geom name="offset_geom" type="sphere" size="0.04"/>
      </body>
    </body>
  </worldbody>
</mujoco>
"""

NUM_ENVS = 4


@pytest.fixture(scope="module")
def device():
  return get_test_device()


class Env:
  def __init__(self, scene, sim, device):
    self.scene = scene
    self.sim = sim
    self.num_envs = scene.num_envs
    self.device = device


def create_test_env(
  device, num_envs=NUM_ENVS, expand_fields=("geom_friction", "dof_damping")
):
  entity_cfg = EntityCfg(spec_fn=lambda: mujoco.MjSpec.from_string(ROBOT_XML))
  scene_cfg = SceneCfg(num_envs=num_envs, entities={"robot": entity_cfg})
  scene = Scene(scene_cfg, device)
  model = scene.compile()

  sim = Simulation(num_envs=num_envs, cfg=SimulationCfg(), model=model, device=device)
  scene.initialize(model, sim.model, sim.data)
  if expand_fields:
    sim.expand_model_fields(expand_fields)

  return Env(scene, sim, device)


@pytest.fixture(scope="module")
def env(device):
  return create_test_env(device)


# Operations: abs, scale, add.


def test_abs_operation(env):
  """Values set directly within range, diversity across envs."""
  torch.manual_seed(42)
  robot = env.scene["robot"]

  dr.geom_friction(
    env,
    env_ids=None,
    ranges=(0.3, 1.2),
    operation="abs",
    asset_cfg=SceneEntityCfg("robot", geom_names=(".*",)),
    axes=[0],
  )

  friction = env.sim.model.geom_friction[:, robot.indexing.geom_ids, 0]
  assert torch.all((friction >= 0.3) & (friction <= 1.2))
  assert len(torch.unique(friction)) >= 2


def test_scale_operation(env):
  """Values = default * sample, within expected bounds."""
  torch.manual_seed(42)
  robot = env.scene["robot"]
  geom_idx = robot.indexing.geom_ids[0]

  default_val = env.sim.get_default_field("geom_friction")[geom_idx, 0].item()

  dr.geom_friction(
    env,
    env_ids=None,
    ranges=(0.5, 2.0),
    operation="scale",
    asset_cfg=SceneEntityCfg("robot", geom_ids=[0]),
    axes=[0],
  )

  result = env.sim.model.geom_friction[:, geom_idx, 0]
  assert torch.all(
    (result >= default_val * 0.5 - 1e-5) & (result <= default_val * 2.0 + 1e-5)
  )


def test_add_operation(env):
  """Values = default + sample, within expected bounds."""
  torch.manual_seed(42)
  robot = env.scene["robot"]
  geom_idx = robot.indexing.geom_ids[0]

  default_val = env.sim.get_default_field("geom_friction")[geom_idx, 0].item()

  dr.geom_friction(
    env,
    env_ids=None,
    ranges=(-0.1, 0.1),
    operation="add",
    asset_cfg=SceneEntityCfg("robot", geom_ids=[0]),
    axes=[0],
  )

  result = env.sim.model.geom_friction[:, geom_idx, 0]
  assert torch.all(
    (result >= default_val - 0.1 - 1e-5) & (result <= default_val + 0.1 + 1e-5)
  )


# No accumulation (scale and add share the same mechanism).


@pytest.mark.parametrize(
  "operation, fixed_val, expected_fn",
  [
    ("scale", (2.0, 2.0), lambda d: d * 2.0),
    ("add", (0.1, 0.1), lambda d: d + 0.1),
  ],
)
def test_no_accumulation(device, operation, fixed_val, expected_fn):
  """3x with fixed value produces same result as 1x (uses defaults)."""
  env = create_test_env(device, num_envs=2)
  robot = env.scene["robot"]
  geom_idx = robot.indexing.geom_ids[0]

  default_friction = env.sim.get_default_field("geom_friction")[geom_idx, 0].item()

  for _ in range(3):
    dr.geom_friction(
      env,
      env_ids=None,
      ranges=fixed_val,
      operation=operation,
      asset_cfg=SceneEntityCfg("robot", geom_ids=[0]),
      axes=[0],
    )

  final = env.sim.model.geom_friction[0, geom_idx, 0].item()
  assert abs(final - expected_fn(default_friction)) < 1e-5


# Distributions.


def test_log_uniform_distribution(device):
  """Log-uniform produces values in [lo, hi], skewed toward lower."""
  torch.manual_seed(42)
  env = create_test_env(device, num_envs=64)

  dr.geom_friction(
    env,
    env_ids=None,
    ranges=(0.1, 10.0),
    operation="abs",
    distribution="log_uniform",
    asset_cfg=SceneEntityCfg("robot", geom_ids=[0]),
    axes=[0],
  )

  vals = env.sim.model.geom_friction[:, env.scene["robot"].indexing.geom_ids[0], 0]
  assert torch.all((vals >= 0.1 - 1e-5) & (vals <= 10.0 + 1e-5))
  # Geometric mean of log-uniform(0.1, 10) is 1.0, well below (0.1+10)/2=5.05.
  assert vals.median() < (0.1 + 10.0) / 2


def test_gaussian_distribution(device):
  """Gaussian: mean/std interpretation, values cluster around mean."""
  torch.manual_seed(42)
  env = create_test_env(device, num_envs=128)

  mean, std = 0.5, 0.05
  dr.geom_friction(
    env,
    env_ids=None,
    ranges=(mean, std),
    operation="abs",
    distribution="gaussian",
    asset_cfg=SceneEntityCfg("robot", geom_ids=[0]),
    axes=[0],
  )

  vals = env.sim.model.geom_friction[:, env.scene["robot"].indexing.geom_ids[0], 0]
  assert abs(vals.mean().item() - mean) < 0.05


# Axes.


@pytest.mark.parametrize(
  "axes, changed, unchanged",
  [
    (None, [0], [1, 2]),  # Default axes for geom_friction = [0]
    ([1, 2], [1, 2], [0]),  # Explicit axes
  ],
)
def test_axes_selectivity(device, axes, changed, unchanged):
  """Only the specified axes are modified; others stay at defaults."""
  env = create_test_env(device, num_envs=2)
  robot = env.scene["robot"]
  geom_idx = robot.indexing.geom_ids[0]

  default_friction = env.sim.get_default_field("geom_friction")[geom_idx].clone()

  dr.geom_friction(
    env,
    env_ids=None,
    ranges=(2.0, 2.0),
    operation="scale",
    asset_cfg=SceneEntityCfg("robot", geom_ids=[0]),
    axes=axes,
  )

  final = env.sim.model.geom_friction[0, geom_idx]
  for ax in changed:
    assert abs(final[ax] - default_friction[ax] * 2.0) < 1e-5
  for ax in unchanged:
    assert abs(final[ax] - default_friction[ax]) < 1e-5


def test_dict_int_axis_ranges(device):
  """{0: (lo, hi), 1: (lo2, hi2)} per-axis ranges."""
  torch.manual_seed(42)
  env = create_test_env(device, num_envs=2)
  robot = env.scene["robot"]
  geom_idx = robot.indexing.geom_ids[0]

  dr.geom_friction(
    env,
    env_ids=None,
    ranges={0: (0.3, 0.4), 1: (0.001, 0.002)},
    operation="abs",
    asset_cfg=SceneEntityCfg("robot", geom_ids=[0]),
  )

  final = env.sim.model.geom_friction[0, geom_idx]
  assert 0.3 - 1e-5 <= final[0].item() <= 0.4 + 1e-5
  assert 0.001 - 1e-5 <= final[1].item() <= 0.002 + 1e-5


def test_invalid_axes_raises(env):
  """axes outside valid_axes raises ValueError."""
  with pytest.raises(ValueError, match="Invalid axes"):
    dr.geom_friction(
      env,
      env_ids=None,
      ranges=(0.3, 1.2),
      operation="abs",
      asset_cfg=SceneEntityCfg("robot", geom_names=(".*",)),
      axes=[3],  # geom_friction valid_axes=[0,1,2]
    )


# String-keyed ranges.


def test_string_keyed_ranges(env):
  """Per-component DR with regex patterns."""
  torch.manual_seed(42)
  dr.joint_damping(
    env,
    env_ids=None,
    ranges={".*joint1": (0.5, 0.5), ".*joint2": (1.5, 1.5)},
    operation="abs",
    asset_cfg=SceneEntityCfg("robot", joint_names=(".*",)),
  )

  robot = env.scene["robot"]
  dof_adr = robot.indexing.joint_v_adr
  damping = env.sim.model.dof_damping[0, dof_adr]
  assert abs(damping[0].item() - 0.5) < 1e-5
  assert abs(damping[1].item() - 1.5) < 1e-5


def test_string_keyed_ranges_no_match_raises(env):
  """Pattern matching no names raises ValueError."""
  with pytest.raises(ValueError, match="matched no"):
    dr.joint_damping(
      env,
      env_ids=None,
      ranges={"nonexistent_joint": (0.5, 1.5)},
      operation="abs",
      asset_cfg=SceneEntityCfg("robot", joint_names=(".*",)),
    )


# Shared random.


def test_shared_random(env):
  """All entities within same env get same value, envs differ."""
  torch.manual_seed(42)
  robot = env.scene["robot"]
  geom_ids = robot.indexing.geom_ids

  dr.geom_friction(
    env,
    env_ids=None,
    ranges=(0.3, 1.2),
    operation="abs",
    asset_cfg=SceneEntityCfg("robot", geom_names=(".*",)),
    axes=[0],
    shared_random=True,
  )

  friction = env.sim.model.geom_friction[:, geom_ids, 0]

  for env_idx in range(env.num_envs):
    env_friction = friction[env_idx]
    assert torch.allclose(env_friction, env_friction[0].expand_as(env_friction))

  env_frictions = friction[:, 0]
  assert len(torch.unique(env_frictions)) > 1
  assert torch.all((friction >= 0.3) & (friction <= 1.2))


# Edge cases.


def test_single_env_without_expand(device):
  """num_envs=1 works without expand_model_fields, no accumulation."""
  env = create_test_env(device, num_envs=1, expand_fields=())
  robot = env.scene["robot"]
  geom_idx = robot.indexing.geom_ids[0]

  original_friction = env.sim.model.geom_friction[0, geom_idx, 0].item()

  for _ in range(2):
    dr.geom_friction(
      env,
      env_ids=None,
      ranges=(2.0, 2.0),
      operation="scale",
      asset_cfg=SceneEntityCfg("robot", geom_ids=[0]),
      axes=[0],
    )

  final = env.sim.model.geom_friction[0, geom_idx, 0].item()
  assert abs(final - original_friction * 2.0) < 1e-5


def test_partial_env_ids(env):
  """Randomizing subset of envs leaves others unchanged."""
  torch.manual_seed(42)
  robot = env.scene["robot"]
  geom_idx = robot.indexing.geom_ids[0]

  original = env.sim.model.geom_friction[:, geom_idx, 0].clone()

  dr.geom_friction(
    env,
    env_ids=torch.tensor([0, 2], device=env.device),
    ranges=(0.1, 0.2),
    operation="abs",
    asset_cfg=SceneEntityCfg("robot", geom_ids=[0]),
    axes=[0],
  )

  result = env.sim.model.geom_friction[:, geom_idx, 0]
  # Envs 0,2 should have changed.
  assert torch.all((result[0] >= 0.1) & (result[0] <= 0.2))
  assert torch.all((result[2] >= 0.1) & (result[2] <= 0.2))
  # Envs 1,3 should be unchanged.
  assert result[1] == original[1]
  assert result[3] == original[3]


# CUDA graph.


@pytest.mark.skipif(
  not torch.cuda.is_available(), reason="CUDA required for graph capture"
)
def test_expand_model_fields_recreates_cuda_graph(device):
  """Verify CUDA graph is recreated after expand_model_fields."""
  entity_cfg = EntityCfg(spec_fn=lambda: mujoco.MjSpec.from_string(ROBOT_XML))
  scene_cfg = SceneCfg(num_envs=NUM_ENVS, entities={"robot": entity_cfg})
  scene = Scene(scene_cfg, device)
  model = scene.compile()

  sim = Simulation(num_envs=NUM_ENVS, cfg=SimulationCfg(), model=model, device=device)
  scene.initialize(model, sim.model, sim.data)

  if not sim.use_cuda_graph:
    pytest.skip("CUDA graph capture not enabled on this device")

  original_step_graph = sim.step_graph

  sim.expand_model_fields(("geom_friction",))

  assert sim.step_graph is not original_step_graph, (
    "CUDA graph was not recreated after expand_model_fields"
  )


# Integration test.


@pytest.mark.slow
def test_g1_foot_friction_shared_across_geoms(device):
  """G1 velocity env has uniform foot friction across all collision geoms."""
  import io
  import warnings
  from contextlib import redirect_stderr, redirect_stdout

  from mjlab.envs.manager_based_rl_env import ManagerBasedRlEnv
  from mjlab.tasks.velocity.config.g1.env_cfgs import unitree_g1_flat_env_cfg

  cfg = unitree_g1_flat_env_cfg()

  with warnings.catch_warnings():
    warnings.simplefilter("ignore")
    with redirect_stdout(io.StringIO()), redirect_stderr(io.StringIO()):
      env = ManagerBasedRlEnv(cfg, device=device)

  try:
    robot = env.scene["robot"]

    foot_geom_names = [
      f"{side}_foot{i}_collision" for side in ("left", "right") for i in range(1, 8)
    ]
    foot_geom_ids, _ = robot.find_geoms(foot_geom_names)
    foot_geom_indices = robot.indexing.geom_ids[foot_geom_ids]

    friction = env.sim.model.geom_friction[:, foot_geom_indices, 0]

    for env_idx in range(env.num_envs):
      env_friction = friction[env_idx]
      assert torch.allclose(env_friction, env_friction[0].expand_as(env_friction))

    if env.num_envs > 1:
      env_frictions = friction[:, 0]
      assert len(torch.unique(env_frictions)) > 1
  finally:
    env.close()


# Quaternion DR tests.


def _make_quat_env(device, num_envs=NUM_ENVS):
  """Create an env with body_quat, geom_quat, site_quat expanded."""
  return create_test_env(
    device,
    num_envs=num_envs,
    expand_fields=("body_quat", "geom_quat", "site_quat"),
  )


@pytest.fixture(scope="module")
def quat_env(device):
  return _make_quat_env(device)


_QUAT_CASES = [
  pytest.param(
    dr.body_quat,
    "body_quat",
    SceneEntityCfg("robot", body_names=(".*",)),
    "body_ids",
    id="body",
  ),
  pytest.param(
    dr.geom_quat,
    "geom_quat",
    SceneEntityCfg("robot", geom_names=(".*",)),
    "geom_ids",
    id="geom",
  ),
  pytest.param(
    dr.site_quat,
    "site_quat",
    SceneEntityCfg("robot", site_names=(".*",)),
    "site_ids",
    id="site",
  ),
]


@pytest.mark.parametrize("dr_func, field, asset_cfg, ids_attr", _QUAT_CASES)
def test_quat_is_unit_quaternion(quat_env, dr_func, field, asset_cfg, ids_attr):
  """All output quaternions must have unit norm (within 1e-6)."""
  torch.manual_seed(0)
  env = quat_env
  robot = env.scene["robot"]

  dr_func(
    env,
    env_ids=None,
    roll_range=(-0.5, 0.5),
    pitch_range=(-0.5, 0.5),
    yaw_range=(-math.pi, math.pi),
    asset_cfg=asset_cfg,
  )

  ids = getattr(robot.indexing, ids_attr)
  quats = getattr(env.sim.model, field)[:, ids, :]
  norms = quats.norm(dim=-1)
  assert torch.all((norms - 1.0).abs() < 1e-6), (
    f"Non-unit quaternion found; max deviation: {(norms - 1.0).abs().max()}"
  )


@pytest.mark.parametrize("dr_func, field, asset_cfg, ids_attr", _QUAT_CASES)
def test_quat_zero_range_unchanged(quat_env, dr_func, field, asset_cfg, ids_attr):
  """All ranges (0, 0): result equals default quaternion."""
  env = quat_env
  robot = env.scene["robot"]
  ids = getattr(robot.indexing, ids_attr)
  default_quat = env.sim.get_default_field(field)[ids].clone()

  dr_func(env, env_ids=None, asset_cfg=asset_cfg)

  result = getattr(env.sim.model, field)[:, ids, :]
  assert torch.allclose(result, default_quat.unsqueeze(0).expand_as(result), atol=1e-6)


@pytest.mark.parametrize("dr_func, field, asset_cfg, ids_attr", _QUAT_CASES)
def test_quat_composes_with_default(quat_env, dr_func, field, asset_cfg, ids_attr):
  """Fixed yaw=pi/4 matches manual quat_mul(q_yaw, q_default)."""
  env = quat_env
  robot = env.scene["robot"]
  ids = getattr(robot.indexing, ids_attr)
  n = len(ids)

  yaw = math.pi / 4
  dr_func(env, env_ids=None, yaw_range=(yaw, yaw), asset_cfg=asset_cfg)

  result = getattr(env.sim.model, field)[:, ids, :]
  q_default = env.sim.get_default_field(field)[ids]
  zeros = torch.zeros(n, device=env.device)
  q_yaw = quat_from_euler_xyz(zeros, zeros, zeros + yaw)
  q_expected = quat_mul(q_yaw, q_default)
  assert torch.allclose(result, q_expected.unsqueeze(0).expand_as(result), atol=1e-6)


def test_body_quat_only_specified_axes(quat_env):
  """Yaw-only perturbation: roll/pitch of result match default."""
  torch.manual_seed(1)
  env = quat_env
  robot = env.scene["robot"]
  body_cfg = SceneEntityCfg("robot", body_names=("base",))
  body_cfg.resolve(env.scene)
  body_ids = robot.indexing.body_ids[body_cfg.body_ids]

  dr.body_quat(
    env,
    env_ids=None,
    yaw_range=(-0.3, 0.3),
    asset_cfg=SceneEntityCfg("robot", body_names=("base",)),
  )

  result = env.sim.model.body_quat[:, body_ids, :]
  # For the default quat [1,0,0,0] composed with a yaw-only rotation, qx and qy should
  # remain 0 (pure yaw -> no roll/pitch).
  q_default = env.sim.get_default_field("body_quat")[body_ids]
  if torch.allclose(q_default, torch.tensor([[1.0, 0.0, 0.0, 0.0]], device=env.device)):
    assert torch.allclose(result[..., 1], torch.zeros_like(result[..., 1]), atol=1e-6)
    assert torch.allclose(result[..., 2], torch.zeros_like(result[..., 2]), atol=1e-6)


# pseudo_inertia tests.


def _make_inertia_env(device, num_envs=NUM_ENVS):
  """Create an env with all inertia-related fields expanded."""
  return create_test_env(
    device,
    num_envs=num_envs,
    expand_fields=("body_mass", "body_ipos", "body_inertia", "body_iquat"),
  )


@pytest.fixture(scope="module")
def inertia_env(device):
  return _make_inertia_env(device)


def _matrix_from_quat(q: torch.Tensor) -> torch.Tensor:
  """Convert wxyz quaternion(s) to 3x3 rotation matrix.

  Shape: (..., 4) → (..., 3, 3).
  """
  w, x, y, z = q[..., 0], q[..., 1], q[..., 2], q[..., 3]
  return torch.stack(
    [
      torch.stack(
        [1 - 2 * (y * y + z * z), 2 * (x * y - z * w), 2 * (x * z + y * w)], dim=-1
      ),
      torch.stack(
        [2 * (x * y + z * w), 1 - 2 * (x * x + z * z), 2 * (y * z - x * w)], dim=-1
      ),
      torch.stack(
        [2 * (x * z - y * w), 2 * (y * z + x * w), 1 - 2 * (x * x + y * y)], dim=-1
      ),
    ],
    dim=-2,
  )


def _reconstruct_J(mass, ipos, inertia, iquat):
  """Reconstruct pseudo-inertia matrix from MuJoCo fields (test helper).

  Correctly accounts for body_iquat (the principal-frame rotation) and the
  parallel-axis theorem (body_ipos) so that the reconstructed J is exact regardless of
  shear magnitude.
  """
  I3 = torch.eye(3, device=mass.device, dtype=mass.dtype)
  # MuJoCo body_iquat maps principal→body.
  R = _matrix_from_quat(iquat)  # (..., 3, 3)
  # Inertia tensor at COM in body frame: I_com = R @ diag(inertia) @ R^T
  I_com = R @ torch.diag_embed(inertia) @ R.mT  # (..., 3, 3)

  # Parallel-axis theorem: shift from COM to body origin.
  c = ipos  # (..., 3)
  c_sq = (c * c).sum(dim=-1)  # (...,)
  c_outer = c.unsqueeze(-1) * c.unsqueeze(-2)  # (..., 3, 3)
  m = mass.unsqueeze(-1).unsqueeze(-1)  # (..., 1, 1)
  I_origin = I_com + m * (c_sq.unsqueeze(-1).unsqueeze(-1) * I3 - c_outer)

  trace = I_origin.diagonal(dim1=-2, dim2=-1).sum(dim=-1)
  sigma = 0.5 * trace.unsqueeze(-1).unsqueeze(-1) * I3 - I_origin
  h = mass.unsqueeze(-1) * ipos
  batch = mass.shape
  J = torch.zeros(*batch, 4, 4, device=mass.device, dtype=mass.dtype)
  J[..., :3, :3] = sigma
  J[..., :3, 3] = h
  J[..., 3, :3] = h
  J[..., 3, 3] = mass
  return J


def test_pseudo_inertia_physical_consistency(inertia_env):
  """After wide perturbation, J ≻ 0, mass > 0, triangle inequality holds."""
  torch.manual_seed(10)
  env = inertia_env
  robot = env.scene["robot"]
  body_cfg = SceneEntityCfg("robot", body_names=("base", "foot1", "foot2"))
  body_cfg.resolve(env.scene)

  dr.pseudo_inertia(
    env,
    env_ids=None,
    alpha_range=(-0.5, 0.5),
    d_range=(-0.3, 0.3),
    s12_range=(-0.1, 0.1),
    s13_range=(-0.1, 0.1),
    s23_range=(-0.1, 0.1),
    t_range=(-0.05, 0.05),
    asset_cfg=body_cfg,
  )

  body_ids = robot.indexing.body_ids[body_cfg.body_ids]
  mass = env.sim.model.body_mass[:, body_ids]
  ipos = env.sim.model.body_ipos[:, body_ids, :]
  inertia = env.sim.model.body_inertia[:, body_ids, :]
  iquat = env.sim.model.body_iquat[:, body_ids, :]

  # Mass must be positive.
  assert torch.all(mass > 0), f"Non-positive mass: {mass.min()}"

  # Principal moments must all be positive.
  assert torch.all(inertia > 0), f"Non-positive principal moment: {inertia.min()}"

  # body_iquat must be a unit quaternion.
  iquat_norms = iquat.norm(dim=-1)
  assert torch.all((iquat_norms - 1.0).abs() < 1e-5), (
    f"Non-unit body_iquat; max deviation: {(iquat_norms - 1.0).abs().max()}"
  )

  # Triangle inequality on principal moments: D_i + D_j >= D_k.
  d1, d2, d3 = inertia[..., 0], inertia[..., 1], inertia[..., 2]
  assert torch.all(d1 + d2 >= d3 - 1e-6)
  assert torch.all(d1 + d3 >= d2 - 1e-6)
  assert torch.all(d2 + d3 >= d1 - 1e-6)

  # Round-trip check: reconstruct J from all four stored fields (including body_iquat)
  # and verify it is positive definite. This is exact because the implementation writes
  # the eigendecomposition of J'.
  J = _reconstruct_J(mass, ipos, inertia, iquat)
  eigvals = torch.linalg.eigvalsh(J)
  assert torch.all(eigvals > -1e-5), f"Non-positive J eigenvalue: {eigvals.min()}"


def test_pseudo_inertia_zero_perturbation_unchanged(inertia_env):
  """All ranges (0, 0): mass, ipos, inertia, iquat equal defaults (1e-5)."""
  env = inertia_env
  robot = env.scene["robot"]
  body_cfg = SceneEntityCfg("robot", body_names=("base", "foot1", "foot2"))
  body_cfg.resolve(env.scene)
  body_ids = robot.indexing.body_ids[body_cfg.body_ids]

  def_mass = env.sim.get_default_field("body_mass")[body_ids].clone()
  def_ipos = env.sim.get_default_field("body_ipos")[body_ids].clone()
  def_inertia = env.sim.get_default_field("body_inertia")[body_ids].clone()
  def_iquat = env.sim.get_default_field("body_iquat")[body_ids].clone()

  dr.pseudo_inertia(env, env_ids=None, asset_cfg=body_cfg)

  mass = env.sim.model.body_mass[:, body_ids]
  ipos = env.sim.model.body_ipos[:, body_ids, :]
  inertia = env.sim.model.body_inertia[:, body_ids, :]
  iquat = env.sim.model.body_iquat[:, body_ids, :]

  assert torch.allclose(mass, def_mass.unsqueeze(0).expand_as(mass), atol=1e-5)
  assert torch.allclose(ipos, def_ipos.unsqueeze(0).expand_as(ipos), atol=1e-5)
  # For bodies with degenerate eigenvalues (2+ equal principal moments), the
  # eigenvector basis is non-unique, so the (inertia, iquat) decomposition may differ
  # while producing the same physical inertia tensor. Compare the full tensor
  # R @ diag(inertia) @ R^T rather than individual components.
  R_def = _matrix_from_quat(def_iquat)
  I_def = R_def @ torch.diag_embed(def_inertia) @ R_def.mT
  R_new = _matrix_from_quat(iquat)
  I_new = R_new @ torch.diag_embed(inertia) @ R_new.mT
  assert torch.allclose(I_new, I_def.unsqueeze(0).expand_as(I_new), atol=1e-5), (
    "Inertia tensor changed unexpectedly under zero perturbation"
  )


def test_pseudo_inertia_zero_perturbation_offset_body(inertia_env):
  """Zero perturbation preserves physics for a body with non-trivial ipos/iquat.

  Since ``eigh`` returns eigenvalues in ascending order while MuJoCo may store them
  differently, we compare the reconstructed pseudo-inertia matrix J (which is
  representation-independent) rather than the raw principal moments and iquat
  individually.
  """
  env = inertia_env
  robot = env.scene["robot"]
  body_cfg = SceneEntityCfg("robot", body_names=("offset_body",))
  body_cfg.resolve(env.scene)
  body_ids = robot.indexing.body_ids[body_cfg.body_ids]

  def_mass = env.sim.get_default_field("body_mass")[body_ids].clone()
  def_ipos = env.sim.get_default_field("body_ipos")[body_ids].clone()
  def_inertia = env.sim.get_default_field("body_inertia")[body_ids].clone()
  def_iquat = env.sim.get_default_field("body_iquat")[body_ids].clone()

  # Verify the test body actually has non-trivial ipos and iquat.
  assert not torch.allclose(def_ipos, torch.zeros_like(def_ipos), atol=1e-6), (
    "Test body should have non-zero ipos"
  )
  identity_q = torch.tensor([1.0, 0.0, 0.0, 0.0], device=env.device)
  assert not (
    torch.allclose(def_iquat, identity_q.expand_as(def_iquat), atol=1e-6)
    or torch.allclose(def_iquat, -identity_q.expand_as(def_iquat), atol=1e-6)
  ), "Test body should have non-identity iquat"

  dr.pseudo_inertia(env, env_ids=None, asset_cfg=body_cfg)

  mass = env.sim.model.body_mass[:, body_ids]
  ipos = env.sim.model.body_ipos[:, body_ids, :]
  inertia = env.sim.model.body_inertia[:, body_ids, :]
  iquat = env.sim.model.body_iquat[:, body_ids, :]

  # Mass and ipos must be exactly preserved.
  assert torch.allclose(mass, def_mass.unsqueeze(0).expand_as(mass), atol=1e-5)
  assert torch.allclose(ipos, def_ipos.unsqueeze(0).expand_as(ipos), atol=1e-5)

  # Inertia and iquat may have different eigenvalue ordering, so compare the
  # reconstructed J matrix instead.
  J_default = _reconstruct_J(def_mass, def_ipos, def_inertia, def_iquat)
  J_result = _reconstruct_J(mass, ipos, inertia, iquat)
  assert torch.allclose(
    J_result, J_default.unsqueeze(0).expand_as(J_result), atol=1e-5
  ), (
    f"Reconstructed J differs under zero perturbation.\n"
    f"Max deviation: {(J_result - J_default.unsqueeze(0)).abs().max()}"
  )


def test_pseudo_inertia_alpha_only_scales_mass(inertia_env):
  """alpha-only: mass and inertia scale by e^{2α}, ipos unchanged."""
  env = inertia_env
  robot = env.scene["robot"]
  body_cfg = SceneEntityCfg("robot", body_names=("base", "foot1", "foot2"))
  body_cfg.resolve(env.scene)
  body_ids = robot.indexing.body_ids[body_cfg.body_ids]

  a = 0.3
  scale = math.exp(2 * a)

  def_mass = env.sim.get_default_field("body_mass")[body_ids].clone()
  def_ipos = env.sim.get_default_field("body_ipos")[body_ids].clone()
  def_inertia = env.sim.get_default_field("body_inertia")[body_ids].clone()

  dr.pseudo_inertia(env, env_ids=None, alpha_range=(a, a), asset_cfg=body_cfg)

  mass = env.sim.model.body_mass[:, body_ids]
  ipos = env.sim.model.body_ipos[:, body_ids, :]
  inertia = env.sim.model.body_inertia[:, body_ids, :]

  expected_mass = def_mass * scale
  expected_inertia = def_inertia * scale

  assert torch.allclose(mass, expected_mass.unsqueeze(0).expand_as(mass), atol=1e-5)
  assert torch.allclose(
    inertia, expected_inertia.unsqueeze(0).expand_as(inertia), atol=1e-5
  )
  assert torch.allclose(ipos, def_ipos.unsqueeze(0).expand_as(ipos), atol=1e-5)


def test_pseudo_inertia_t1_shifts_com(inertia_env):
  """t1-only: mass unchanged, ipos[0] shifts by exactly t1."""
  env = inertia_env
  robot = env.scene["robot"]
  body_cfg = SceneEntityCfg("robot", body_names=("base", "foot1", "foot2"))
  body_cfg.resolve(env.scene)
  body_ids = robot.indexing.body_ids[body_cfg.body_ids]

  t = 0.05
  def_mass = env.sim.get_default_field("body_mass")[body_ids].clone()
  def_ipos = env.sim.get_default_field("body_ipos")[body_ids].clone()

  dr.pseudo_inertia(env, env_ids=None, t1_range=(t, t), asset_cfg=body_cfg)

  mass = env.sim.model.body_mass[:, body_ids]
  ipos = env.sim.model.body_ipos[:, body_ids, :]

  assert torch.allclose(mass, def_mass.unsqueeze(0).expand_as(mass), atol=1e-5)
  # ipos x-component shifts by t; y and z are unchanged.
  assert torch.allclose(
    ipos[..., 0], (def_ipos[..., 0] + t).unsqueeze(0).expand_as(ipos[..., 0]), atol=1e-5
  )
  assert torch.allclose(
    ipos[..., 1:], def_ipos[..., 1:].unsqueeze(0).expand_as(ipos[..., 1:]), atol=1e-5
  )


def test_pseudo_inertia_d_isotropic_vs_per_axis(device):
  """d_range=(r,r) gives same result as d1=d2=d3=(r,r)."""
  r = 0.2
  seed = 42

  torch.manual_seed(seed)
  env1 = _make_inertia_env(device)
  body_cfg = SceneEntityCfg("robot", body_names=("base", "foot1", "foot2"))
  body_cfg.resolve(env1.scene)
  dr.pseudo_inertia(env1, env_ids=None, d_range=(r, r), asset_cfg=body_cfg)

  torch.manual_seed(seed)
  env2 = _make_inertia_env(device)
  dr.pseudo_inertia(
    env2,
    env_ids=None,
    d1_range=(r, r),
    d2_range=(r, r),
    d3_range=(r, r),
    asset_cfg=body_cfg,
  )

  robot1 = env1.scene["robot"]
  body_ids = robot1.indexing.body_ids[body_cfg.body_ids]

  assert torch.allclose(
    env1.sim.model.body_mass[:, body_ids],
    env2.sim.model.body_mass[:, body_ids],
    atol=1e-6,
  )
  assert torch.allclose(
    env1.sim.model.body_inertia[:, body_ids, :],
    env2.sim.model.body_inertia[:, body_ids, :],
    atol=1e-6,
  )


def test_pseudo_inertia_no_accumulation(device):
  """3x with alpha=0.1 gives same result as 1x (uses defaults each time)."""
  env = _make_inertia_env(device)
  robot = env.scene["robot"]
  body_cfg = SceneEntityCfg("robot", body_names=("base",))
  body_cfg.resolve(env.scene)
  body_ids = robot.indexing.body_ids[body_cfg.body_ids]

  for _ in range(3):
    dr.pseudo_inertia(env, env_ids=None, alpha_range=(0.1, 0.1), asset_cfg=body_cfg)

  mass_3x = env.sim.model.body_mass[:, body_ids].clone()

  env2 = _make_inertia_env(device)
  dr.pseudo_inertia(env2, env_ids=None, alpha_range=(0.1, 0.1), asset_cfg=body_cfg)
  mass_1x = env2.sim.model.body_mass[:, body_ids]

  assert torch.allclose(mass_3x, mass_1x, atol=1e-5)


def test_pseudo_inertia_partial_env_ids(device):
  """Randomizing only env 0 leaves env 1 unchanged."""
  env = _make_inertia_env(device, num_envs=2)
  robot = env.scene["robot"]
  body_cfg = SceneEntityCfg("robot", body_names=("base",))
  body_cfg.resolve(env.scene)
  body_ids = robot.indexing.body_ids[body_cfg.body_ids]

  mass_before = env.sim.model.body_mass[:, body_ids].clone()

  dr.pseudo_inertia(
    env,
    env_ids=torch.tensor([0], device=device),
    alpha_range=(0.5, 0.5),
    asset_cfg=body_cfg,
  )

  mass_after = env.sim.model.body_mass[:, body_ids]
  # Env 0 changed.
  assert not torch.allclose(mass_after[0], mass_before[0], atol=1e-6)
  # Env 1 unchanged.
  assert torch.allclose(mass_after[1], mass_before[1], atol=1e-6)


# Camera / Light DR tests.


def _make_cam_light_env(device, num_envs=NUM_ENVS):
  """Create an env with camera and light fields expanded."""
  return create_test_env(
    device,
    num_envs=num_envs,
    expand_fields=(
      "cam_fovy",
      "cam_pos",
      "cam_quat",
      "cam_intrinsic",
      "light_pos",
      "light_dir",
      "light_diffuse",
      "light_specular",
      "light_ambient",
      "light_attenuation",
      "light_cutoff",
      "light_exponent",
    ),
  )


@pytest.fixture(scope="module")
def cam_light_env(device):
  return _make_cam_light_env(device)


def test_cam_fovy_abs(cam_light_env):
  """Set fovy to absolute range, check bounds."""
  torch.manual_seed(42)
  env = cam_light_env
  robot = env.scene["robot"]
  cam_cfg = SceneEntityCfg("robot", camera_names=(".*",))
  cam_cfg.resolve(env.scene)
  cam_ids = robot.indexing.cam_ids[cam_cfg.camera_ids]

  dr.cam_fovy(
    env,
    env_ids=None,
    ranges=(30.0, 90.0),
    operation="abs",
    asset_cfg=cam_cfg,
  )

  fovy = env.sim.model.cam_fovy[:, cam_ids]
  assert torch.all((fovy >= 30.0 - 1e-3) & (fovy <= 90.0 + 1e-3))


def test_cam_pos_add(cam_light_env):
  """Add offset to camera positions, check default + offset."""
  torch.manual_seed(42)
  env = cam_light_env
  robot = env.scene["robot"]
  cam_cfg = SceneEntityCfg("robot", camera_names=("front_cam",))
  cam_cfg.resolve(env.scene)
  cam_ids = robot.indexing.cam_ids[cam_cfg.camera_ids]

  default_pos = env.sim.get_default_field("cam_pos")[cam_ids].clone()

  dr.cam_pos(
    env,
    env_ids=None,
    ranges=(-0.1, 0.1),
    operation="add",
    asset_cfg=cam_cfg,
  )

  result = env.sim.model.cam_pos[:, cam_ids, :]
  for ax in range(3):
    lo = default_pos[..., ax] - 0.1 - 1e-5
    hi = default_pos[..., ax] + 0.1 + 1e-5
    assert torch.all((result[..., ax] >= lo) & (result[..., ax] <= hi))


def test_cam_quat_zero_unchanged(cam_light_env):
  """Zero RPY range preserves default quaternion."""
  env = cam_light_env
  robot = env.scene["robot"]
  cam_cfg = SceneEntityCfg("robot", camera_names=(".*",))
  cam_cfg.resolve(env.scene)
  cam_ids = robot.indexing.cam_ids[cam_cfg.camera_ids]

  default_quat = env.sim.get_default_field("cam_quat")[cam_ids].clone()

  dr.cam_quat(env, env_ids=None, asset_cfg=cam_cfg)

  result = env.sim.model.cam_quat[:, cam_ids, :]
  assert torch.allclose(
    result,
    default_quat.unsqueeze(0).expand_as(result),
    atol=1e-6,
  )


def test_light_pos_abs(cam_light_env):
  """Set light position to absolute range, check bounds."""
  torch.manual_seed(42)
  env = cam_light_env
  robot = env.scene["robot"]
  light_cfg = SceneEntityCfg("robot", light_names=(".*",))
  light_cfg.resolve(env.scene)
  light_ids = robot.indexing.light_ids[light_cfg.light_ids]

  dr.light_pos(
    env,
    env_ids=None,
    ranges=(-5.0, 5.0),
    operation="abs",
    asset_cfg=light_cfg,
  )

  pos = env.sim.model.light_pos[:, light_ids, :]
  assert torch.all((pos >= -5.0 - 1e-3) & (pos <= 5.0 + 1e-3))


def test_light_dir_add(cam_light_env):
  """Add offset to light direction, check default + offset."""
  torch.manual_seed(42)
  env = cam_light_env
  robot = env.scene["robot"]
  light_cfg = SceneEntityCfg("robot", light_names=(".*",))
  light_cfg.resolve(env.scene)
  light_ids = robot.indexing.light_ids[light_cfg.light_ids]

  default_dir = env.sim.get_default_field("light_dir")[light_ids].clone()

  dr.light_dir(
    env,
    env_ids=None,
    ranges=(-0.5, 0.5),
    operation="add",
    asset_cfg=light_cfg,
  )

  result = env.sim.model.light_dir[:, light_ids, :]
  for ax in range(3):
    lo = default_dir[..., ax] - 0.5 - 1e-5
    hi = default_dir[..., ax] + 0.5 + 1e-5
    assert torch.all((result[..., ax] >= lo) & (result[..., ax] <= hi))


@pytest.mark.parametrize(
  ("func_name", "field", "ranges"),
  [
    ("light_diffuse", "light_diffuse", (0.2, 0.8)),
    ("light_specular", "light_specular", (0.0, 1.0)),
    ("light_ambient", "light_ambient", (0.1, 0.5)),
    ("light_attenuation", "light_attenuation", (0.0, 1.0)),
  ],
)
def test_light_vec3_fields_abs(cam_light_env, func_name, field, ranges):
  """Vec3 light fields randomize per selected light."""
  torch.manual_seed(42)
  env = cam_light_env
  robot = env.scene["robot"]
  light_cfg = SceneEntityCfg("robot", light_names=(".*",))
  light_cfg.resolve(env.scene)
  light_ids = robot.indexing.light_ids[light_cfg.light_ids]
  func = getattr(dr, func_name)

  func(
    env,
    env_ids=None,
    ranges=ranges,
    operation="abs",
    asset_cfg=light_cfg,
  )

  values = getattr(env.sim.model, field)[:, light_ids, :]
  lower, upper = ranges
  assert torch.all((values >= lower - 1e-5) & (values <= upper + 1e-5))
  assert len(torch.unique(values[:, 0, 0])) >= 2


@pytest.mark.parametrize(
  ("func_name", "field", "ranges"),
  [
    ("light_cutoff", "light_cutoff", (20.0, 60.0)),
    ("light_exponent", "light_exponent", (1.0, 20.0)),
  ],
)
def test_light_scalar_fields_abs(cam_light_env, func_name, field, ranges):
  """Scalar light fields randomize per selected light."""
  torch.manual_seed(42)
  env = cam_light_env
  robot = env.scene["robot"]
  light_cfg = SceneEntityCfg("robot", light_names=(".*",))
  light_cfg.resolve(env.scene)
  light_ids = robot.indexing.light_ids[light_cfg.light_ids]
  func = getattr(dr, func_name)

  func(
    env,
    env_ids=None,
    ranges=ranges,
    operation="abs",
    asset_cfg=light_cfg,
  )

  values = getattr(env.sim.model, field)[:, light_ids]
  lower, upper = ranges
  assert torch.all((values >= lower - 1e-5) & (values <= upper + 1e-5))
  assert len(torch.unique(values[:, 0])) >= 2


def test_camera_partial_env_ids(cam_light_env):
  """Subset of envs randomized, others unchanged."""
  torch.manual_seed(42)
  env = cam_light_env
  robot = env.scene["robot"]
  cam_cfg = SceneEntityCfg("robot", camera_names=(".*",))
  cam_cfg.resolve(env.scene)
  cam_ids = robot.indexing.cam_ids[cam_cfg.camera_ids]

  original_fovy = env.sim.model.cam_fovy[:, cam_ids].clone()

  dr.cam_fovy(
    env,
    env_ids=torch.tensor([0, 2], device=env.device),
    ranges=(10.0, 20.0),
    operation="abs",
    asset_cfg=cam_cfg,
  )

  result = env.sim.model.cam_fovy[:, cam_ids]
  # Envs 0, 2 should have changed.
  assert torch.all((result[0] >= 10.0 - 1e-3) & (result[0] <= 20.0 + 1e-3))
  assert torch.all((result[2] >= 10.0 - 1e-3) & (result[2] <= 20.0 + 1e-3))
  # Envs 1, 3 should be unchanged.
  assert torch.allclose(result[1], original_fovy[1])
  assert torch.allclose(result[3], original_fovy[3])


# geom_size DR tests.

GEOM_SIZE_XML = """
<mujoco>
  <worldbody>
    <body name="base" pos="0 0 1">
      <freejoint name="free_joint"/>
      <geom name="box_geom" type="box" size="0.1 0.2 0.3" mass="1.0"/>
      <geom name="sphere_geom" type="sphere" size="0.15" mass="0.5"/>
      <geom name="capsule_geom" type="capsule" size="0.05 0.2" mass="0.3"/>
      <geom name="cylinder_geom" type="cylinder" size="0.08 0.15" mass="0.4"/>
      <geom name="ellipsoid_geom" type="ellipsoid" size="0.1 0.2 0.05"
        mass="0.2"/>
    </body>
  </worldbody>
</mujoco>
"""


def _make_geom_size_env(device, num_envs=NUM_ENVS):
  entity_cfg = EntityCfg(spec_fn=lambda: mujoco.MjSpec.from_string(GEOM_SIZE_XML))
  scene_cfg = SceneCfg(num_envs=num_envs, entities={"robot": entity_cfg})
  scene = Scene(scene_cfg, device)
  model = scene.compile()
  sim = Simulation(num_envs=num_envs, cfg=SimulationCfg(), model=model, device=device)
  scene.initialize(model, sim.model, sim.data)
  sim.expand_model_fields(("geom_size", "geom_rbound", "geom_aabb"))
  return Env(scene, sim, device)


@pytest.fixture(scope="module")
def geom_size_env(device):
  return _make_geom_size_env(device)


def _expected_rbound(geom_type, s0, s1, s2):
  """Reference rbound computation matching MuJoCo GetRBound."""
  if geom_type == "sphere":
    return s0
  elif geom_type == "capsule":
    return s0 + s1
  elif geom_type == "cylinder":
    return math.sqrt(s0**2 + s1**2)
  elif geom_type == "ellipsoid":
    return max(s0, s1, s2)
  elif geom_type == "box":
    return math.sqrt(s0**2 + s1**2 + s2**2)
  raise ValueError(geom_type)


def _expected_aabb_half(geom_type, s0, s1, s2):
  """Reference aabb half-size computation matching MuJoCo ComputeAABB."""
  if geom_type == "sphere":
    return (s0, s0, s0)
  elif geom_type == "capsule":
    return (s0, s0, s0 + s1)
  elif geom_type == "cylinder":
    return (s0, s0, s1)
  elif geom_type == "ellipsoid":
    return (s0, s1, s2)
  elif geom_type == "box":
    return (s0, s1, s2)
  raise ValueError(geom_type)


def test_geom_size_scale_updates_bounds(geom_size_env):
  """Scaling geom_size updates rbound and aabb consistently."""
  torch.manual_seed(42)
  env = geom_size_env
  robot = env.scene["robot"]
  geom_cfg = SceneEntityCfg("robot", geom_names=(".*",))
  geom_cfg.resolve(env.scene)

  dr.geom_size(
    env,
    env_ids=None,
    ranges=(0.8, 1.5),
    operation="scale",
    asset_cfg=geom_cfg,
  )

  geom_ids = robot.indexing.geom_ids[geom_cfg.geom_ids]
  size = env.sim.model.geom_size[:, geom_ids]  # (E, G, 3)
  rbound = env.sim.model.geom_rbound[:, geom_ids]  # (E, G)
  aabb = env.sim.model.geom_aabb[:, geom_ids]  # (E, G, 2, 3)

  geom_names = list(robot.geom_names)
  type_names = {
    "box_geom": "box",
    "sphere_geom": "sphere",
    "capsule_geom": "capsule",
    "cylinder_geom": "cylinder",
    "ellipsoid_geom": "ellipsoid",
  }

  for g_local, gname in enumerate(geom_names):
    tname = type_names[gname]
    for e in range(env.num_envs):
      s0 = size[e, g_local, 0].item()
      s1 = size[e, g_local, 1].item()
      s2 = size[e, g_local, 2].item()

      expected_rb = _expected_rbound(tname, s0, s1, s2)
      actual_rb = rbound[e, g_local].item()
      assert abs(actual_rb - expected_rb) < 1e-5, (
        f"{gname} env {e}: rbound {actual_rb} != {expected_rb}"
      )

      ex, ey, ez = _expected_aabb_half(tname, s0, s1, s2)
      actual_half = aabb[e, g_local, 1]  # half-size is index 1
      assert abs(actual_half[0].item() - ex) < 1e-5
      assert abs(actual_half[1].item() - ey) < 1e-5
      assert abs(actual_half[2].item() - ez) < 1e-5

      # Center should be zero for all primitives.
      center = aabb[e, g_local, 0]
      assert torch.allclose(center, torch.zeros(3, device=env.device), atol=1e-6)


def test_geom_size_no_accumulation(device):
  """Repeated scale operations don't accumulate."""
  env = _make_geom_size_env(device, num_envs=2)
  robot = env.scene["robot"]
  geom_cfg = SceneEntityCfg("robot", geom_names=("box_geom",))
  geom_cfg.resolve(env.scene)
  geom_ids = robot.indexing.geom_ids[geom_cfg.geom_ids]

  default_size = env.sim.get_default_field("geom_size")[geom_ids].clone()

  for _ in range(3):
    dr.geom_size(
      env,
      env_ids=None,
      ranges=(2.0, 2.0),
      operation="scale",
      asset_cfg=geom_cfg,
    )

  result = env.sim.model.geom_size[0, geom_ids]
  expected = default_size * 2.0
  assert torch.allclose(result, expected, atol=1e-5)


def test_geom_size_separate_axes_dont_clobber(device):
  """Events targeting different axes of the same geom compose, not clobber."""
  env = _make_geom_size_env(device, num_envs=2)
  robot = env.scene["robot"]
  geom_cfg = SceneEntityCfg("robot", geom_names=("box_geom",))
  geom_cfg.resolve(env.scene)
  geom_ids = robot.indexing.geom_ids[geom_cfg.geom_ids]

  default_size = env.sim.get_default_field("geom_size")[geom_ids].clone()

  # Scale axis 0 by 0.1, then axis 1 by 3.0 in a separate event.
  dr.geom_size(
    env,
    env_ids=None,
    ranges=(0.1, 0.1),
    operation="scale",
    axes=[0],
    asset_cfg=geom_cfg,
  )
  dr.geom_size(
    env,
    env_ids=None,
    ranges=(3.0, 3.0),
    operation="scale",
    axes=[1],
    asset_cfg=geom_cfg,
  )

  result = env.sim.model.geom_size[0, geom_ids]
  expected = default_size.clone()
  expected[..., 0] *= 0.1
  expected[..., 1] *= 3.0
  assert torch.allclose(result, expected, atol=1e-5)


def test_geom_size_partial_env_ids(device):
  """Only specified envs are updated."""
  env = _make_geom_size_env(device, num_envs=4)
  robot = env.scene["robot"]
  geom_cfg = SceneEntityCfg("robot", geom_names=("sphere_geom",))
  geom_cfg.resolve(env.scene)
  geom_ids = robot.indexing.geom_ids[geom_cfg.geom_ids]

  original_rbound = env.sim.model.geom_rbound[:, geom_ids].clone()

  dr.geom_size(
    env,
    env_ids=torch.tensor([0, 2], device=device),
    ranges=(2.0, 2.0),
    operation="scale",
    asset_cfg=geom_cfg,
  )

  rbound = env.sim.model.geom_rbound[:, geom_ids]
  # Envs 0, 2 should have changed.
  assert not torch.allclose(rbound[0], original_rbound[0])
  assert not torch.allclose(rbound[2], original_rbound[2])
  # Envs 1, 3 should be unchanged.
  assert torch.allclose(rbound[1], original_rbound[1])
  assert torch.allclose(rbound[3], original_rbound[3])


def test_geom_size_raises_on_unsupported_type(device):
  """dr.geom_size raises ValueError when a non-primitive geom type is selected."""
  PLANE_XML = """
<mujoco>
  <worldbody>
    <geom name="floor" type="plane" size="5 5 0.1"/>
    <body name="base" pos="0 0 1">
      <freejoint name="free_joint"/>
      <geom name="box_geom" type="box" size="0.1 0.1 0.1" mass="1.0"/>
    </body>
  </worldbody>
</mujoco>
"""
  entity_cfg = EntityCfg(spec_fn=lambda: mujoco.MjSpec.from_string(PLANE_XML))
  scene_cfg = SceneCfg(num_envs=2, entities={"robot": entity_cfg})
  scene = Scene(scene_cfg, device)
  model = scene.compile()
  sim = Simulation(num_envs=2, cfg=SimulationCfg(), model=model, device=device)
  scene.initialize(model, sim.model, sim.data)
  sim.expand_model_fields(("geom_size", "geom_rbound", "geom_aabb"))
  env = Env(scene, sim, device)

  plane_cfg = SceneEntityCfg("robot", geom_names=("floor",))
  plane_cfg.resolve(env.scene)

  with pytest.raises(ValueError, match="unsupported types"):
    dr.geom_size(env, env_ids=None, ranges=(1.0, 2.0), asset_cfg=plane_cfg)


# Tendon DR tests.

TENDON_XML = """
<mujoco>
  <worldbody>
    <body name="base" pos="0 0 1">
      <freejoint name="free_joint"/>
      <geom name="base_geom" type="box" size="0.1 0.1 0.1" mass="1.0"/>
      <site name="anchor" pos="0 0 0.1"/>
      <body name="child" pos="0.3 0 0">
        <joint name="hinge1" type="hinge" axis="0 0 1"/>
        <geom name="child_geom" type="sphere" size="0.05" mass="0.2"/>
        <site name="hook" pos="0 0 -0.05"/>
      </body>
    </body>
  </worldbody>
  <tendon>
    <spatial name="wire" limited="true" range="0 0.5" width="0.003"
      damping="2.0" stiffness="10.0" frictionloss="0.5" armature="0.1">
      <site site="anchor"/>
      <site site="hook"/>
    </spatial>
  </tendon>
</mujoco>
"""


def _make_tendon_env(device, num_envs=NUM_ENVS):
  entity_cfg = EntityCfg(spec_fn=lambda: mujoco.MjSpec.from_string(TENDON_XML))
  scene_cfg = SceneCfg(num_envs=num_envs, entities={"robot": entity_cfg})
  scene = Scene(scene_cfg, device)
  model = scene.compile()
  sim = Simulation(num_envs=num_envs, cfg=SimulationCfg(), model=model, device=device)
  scene.initialize(model, sim.model, sim.data)
  sim.expand_model_fields(
    (
      "tendon_damping",
      "tendon_stiffness",
      "tendon_frictionloss",
      "tendon_armature",
    )
  )

  return Env(scene, sim, device)


@pytest.fixture(scope="module")
def tendon_env(device):
  return _make_tendon_env(device)


@pytest.mark.parametrize(
  "dr_func, field",
  [
    (dr.tendon_damping, "tendon_damping"),
    (dr.tendon_stiffness, "tendon_stiffness"),
    (dr.tendon_friction, "tendon_frictionloss"),
    (dr.tendon_armature, "tendon_armature"),
  ],
)
def test_tendon_field_scale(tendon_env, dr_func, field):
  """Scale operation on tendon fields: result = default * sample."""
  torch.manual_seed(42)
  env = tendon_env
  robot = env.scene["robot"]
  tendon_cfg = SceneEntityCfg("robot", tendon_names=(".*",))
  tendon_cfg.resolve(env.scene)
  tendon_ids = robot.indexing.tendon_ids[tendon_cfg.tendon_ids]

  default_val = env.sim.get_default_field(field)[tendon_ids].clone()

  dr_func(
    env,
    env_ids=None,
    ranges=(0.5, 2.0),
    operation="scale",
    asset_cfg=tendon_cfg,
  )

  result = getattr(env.sim.model, field)[:, tendon_ids]
  lo = default_val * 0.5 - 1e-5
  hi = default_val * 2.0 + 1e-5
  assert torch.all((result >= lo) & (result <= hi))


def test_tendon_armature_no_accumulation(device):
  """Repeated scale does not accumulate."""
  env = _make_tendon_env(device, num_envs=2)
  robot = env.scene["robot"]
  tendon_cfg = SceneEntityCfg("robot", tendon_names=(".*",))
  tendon_cfg.resolve(env.scene)
  tendon_ids = robot.indexing.tendon_ids[tendon_cfg.tendon_ids]

  default_val = env.sim.get_default_field("tendon_armature")[tendon_ids]

  for _ in range(3):
    dr.tendon_armature(
      env,
      env_ids=None,
      ranges=(2.0, 2.0),
      operation="scale",
      asset_cfg=tendon_cfg,
    )

  result = env.sim.model.tendon_armature[0, tendon_ids]
  assert torch.allclose(result, default_val * 2.0, atol=1e-5)


# Extensible Operation / Distribution types.


def test_operation_instance_abs(device):
  """Passing an Operation instance works identically to the string."""
  from mjlab.envs.mdp.dr._types import abs as abs_op

  torch.manual_seed(42)
  env1 = create_test_env(device)
  robot = env1.scene["robot"]

  dr.geom_friction(
    env1,
    env_ids=None,
    ranges=(0.3, 1.2),
    operation="abs",
    asset_cfg=SceneEntityCfg("robot", geom_names=(".*",)),
    axes=[0],
  )
  result_str = env1.sim.model.geom_friction[:, robot.indexing.geom_ids, 0].clone()

  torch.manual_seed(42)
  env2 = create_test_env(device)

  dr.geom_friction(
    env2,
    env_ids=None,
    ranges=(0.3, 1.2),
    operation=abs_op,
    asset_cfg=SceneEntityCfg("robot", geom_names=(".*",)),
    axes=[0],
  )
  result_inst = env2.sim.model.geom_friction[:, robot.indexing.geom_ids, 0]

  assert torch.allclose(result_str, result_inst)


def test_distribution_instance_uniform(device):
  """Passing a Distribution instance works identically to the string."""
  from mjlab.envs.mdp.dr._types import uniform as uniform_dist

  torch.manual_seed(42)
  env1 = create_test_env(device)
  robot = env1.scene["robot"]

  dr.geom_friction(
    env1,
    env_ids=None,
    ranges=(0.3, 1.2),
    distribution="uniform",
    operation="abs",
    asset_cfg=SceneEntityCfg("robot", geom_names=(".*",)),
    axes=[0],
  )
  result_str = env1.sim.model.geom_friction[:, robot.indexing.geom_ids, 0].clone()

  torch.manual_seed(42)
  env2 = create_test_env(device)

  dr.geom_friction(
    env2,
    env_ids=None,
    ranges=(0.3, 1.2),
    distribution=uniform_dist,
    operation="abs",
    asset_cfg=SceneEntityCfg("robot", geom_names=(".*",)),
    axes=[0],
  )
  result_inst = env2.sim.model.geom_friction[:, robot.indexing.geom_ids, 0]

  assert torch.allclose(result_str, result_inst)


def test_custom_operation(env):
  """A user-defined Operation works end-to-end."""
  from mjlab.envs.mdp.dr._types import Operation

  clamp_op = Operation(
    name="clamp",
    initialize=torch.Tensor.clone,
    combine=lambda _base, random: torch.clamp(random, min=0.4, max=0.9),
    uses_defaults=False,
  )

  torch.manual_seed(42)
  robot = env.scene["robot"]

  dr.geom_friction(
    env,
    env_ids=None,
    ranges=(0.1, 1.5),
    operation=clamp_op,
    asset_cfg=SceneEntityCfg("robot", geom_names=(".*",)),
    axes=[0],
  )

  friction = env.sim.model.geom_friction[:, robot.indexing.geom_ids, 0]
  assert torch.all((friction >= 0.4 - 1e-5) & (friction <= 0.9 + 1e-5))


def test_custom_distribution(env):
  """A user-defined Distribution works end-to-end."""
  from mjlab.envs.mdp.dr._types import Distribution

  # Distribution that always returns the midpoint.
  midpoint_dist = Distribution(
    name="midpoint",
    sample=lambda lo, hi, shape, _device: ((lo + hi) / 2).expand(shape),
  )

  robot = env.scene["robot"]

  dr.geom_friction(
    env,
    env_ids=None,
    ranges=(0.2, 0.8),
    distribution=midpoint_dist,
    operation="abs",
    asset_cfg=SceneEntityCfg("robot", geom_names=(".*",)),
    axes=[0],
  )

  friction = env.sim.model.geom_friction[:, robot.indexing.geom_ids, 0]
  assert torch.allclose(friction, torch.tensor(0.5, device=env.device), atol=1e-5)


def test_resolve_unknown_operation_raises(device):
  """Unknown operation string raises ValueError."""
  from mjlab.envs.mdp.dr._types import resolve_operation

  with pytest.raises(ValueError, match="Unknown operation"):
    resolve_operation("nonexistent")


def test_resolve_unknown_distribution_raises(device):
  """Unknown distribution string raises ValueError."""
  from mjlab.envs.mdp.dr._types import resolve_distribution

  with pytest.raises(ValueError, match="Unknown distribution"):
    resolve_distribution("nonexistent")


# mat_rgba tests.

MAT_XML = """
<mujoco>
  <asset>
    <material name="test_mat" rgba="1 1 1 1" emission="0.1"
      specular="0.2" shininess="0.3" texrepeat="1 2"/>
  </asset>
  <worldbody>
    <body name="base" pos="0 0 1">
      <freejoint name="free_joint"/>
      <geom name="base_geom" type="box" size="0.1 0.1 0.1" mass="1.0"
        material="test_mat"/>
    </body>
  </worldbody>
</mujoco>
"""


def _make_mat_env(device, num_envs=NUM_ENVS):
  entity_cfg = EntityCfg(spec_fn=lambda: mujoco.MjSpec.from_string(MAT_XML))
  scene_cfg = SceneCfg(num_envs=num_envs, entities={"robot": entity_cfg})
  scene = Scene(scene_cfg, device)
  model = scene.compile()
  sim = Simulation(num_envs=num_envs, cfg=SimulationCfg(), model=model, device=device)
  scene.initialize(model, sim.model, sim.data)
  sim.expand_model_fields(
    ("mat_rgba", "mat_emission", "mat_specular", "mat_shininess", "mat_texrepeat")
  )
  return Env(scene, sim, device)


@pytest.fixture(scope="module")
def mat_env(device):
  return _make_mat_env(device)


def test_mat_rgba_abs(mat_env):
  """Values within range, diversity across envs."""
  torch.manual_seed(42)
  env = mat_env

  dr.mat_rgba(
    env,
    env_ids=None,
    ranges=(0.2, 0.8),
    asset_cfg=SceneEntityCfg("robot", material_names=("test_mat",)),
    operation="abs",
  )

  mat_id = mujoco.mj_name2id(
    env.sim.mj_model, mujoco.mjtObj.mjOBJ_MATERIAL, "robot/test_mat"
  )
  rgba = env.sim.model.mat_rgba[:, mat_id, :]
  assert torch.all((rgba >= 0.2 - 1e-5) & (rgba <= 0.8 + 1e-5))
  assert len(torch.unique(rgba[:, 0])) >= 2


@pytest.mark.parametrize(
  ("func_name", "field", "ranges"),
  [
    ("mat_emission", "mat_emission", (0.1, 0.7)),
    ("mat_specular", "mat_specular", (0.0, 1.0)),
    ("mat_shininess", "mat_shininess", (0.2, 0.9)),
  ],
)
def test_scalar_material_fields_abs(mat_env, func_name, field, ranges):
  """Scalar material fields randomize per selected material."""
  torch.manual_seed(42)
  env = mat_env
  func = getattr(dr, func_name)

  func(
    env,
    env_ids=None,
    ranges=ranges,
    asset_cfg=SceneEntityCfg("robot", material_names=("test_mat",)),
    operation="abs",
  )

  mat_id = mujoco.mj_name2id(
    env.sim.mj_model, mujoco.mjtObj.mjOBJ_MATERIAL, "robot/test_mat"
  )
  values = getattr(env.sim.model, field)[:, mat_id]
  lower, upper = ranges
  assert torch.all((values >= lower - 1e-5) & (values <= upper + 1e-5))
  assert len(torch.unique(values)) >= 2


def test_mat_texrepeat_abs(mat_env):
  """Texture repeat randomizes S/T axes for selected materials."""
  torch.manual_seed(42)
  env = mat_env

  dr.mat_texrepeat(
    env,
    env_ids=None,
    ranges={0: (1.0, 3.0), 1: (4.0, 6.0)},
    asset_cfg=SceneEntityCfg("robot", material_names=("test_mat",)),
    operation="abs",
  )

  mat_id = mujoco.mj_name2id(
    env.sim.mj_model, mujoco.mjtObj.mjOBJ_MATERIAL, "robot/test_mat"
  )
  texrepeat = env.sim.model.mat_texrepeat[:, mat_id, :]
  assert torch.all((texrepeat[:, 0] >= 1.0 - 1e-5) & (texrepeat[:, 0] <= 3.0 + 1e-5))
  assert torch.all((texrepeat[:, 1] >= 4.0 - 1e-5) & (texrepeat[:, 1] <= 6.0 + 1e-5))
  assert len(torch.unique(texrepeat[:, 0])) >= 2


@pytest.mark.parametrize(
  ("func_name", "field"),
  [
    ("mat_emission", "mat_emission"),
    ("mat_specular", "mat_specular"),
    ("mat_shininess", "mat_shininess"),
  ],
)
def test_scalar_material_fields_scale_use_defaults(mat_env, func_name, field):
  """Scale operation uses compiled defaults for scalar material fields."""
  env = mat_env
  func = getattr(dr, func_name)

  func(
    env,
    env_ids=None,
    ranges=(2.0, 2.0),
    asset_cfg=SceneEntityCfg("robot", material_names=("test_mat",)),
    operation="scale",
  )

  mat_id = mujoco.mj_name2id(
    env.sim.mj_model, mujoco.mjtObj.mjOBJ_MATERIAL, "robot/test_mat"
  )
  expected = env.sim.get_default_field(field)[mat_id] * 2.0
  values = getattr(env.sim.model, field)[:, mat_id]
  assert torch.allclose(values, expected.expand_as(values))


def test_mat_texrepeat_scale_uses_defaults(mat_env):
  """Scale operation uses compiled defaults and both S/T axes by default."""
  env = mat_env

  dr.mat_texrepeat(
    env,
    env_ids=None,
    ranges=(2.0, 2.0),
    asset_cfg=SceneEntityCfg("robot", material_names=("test_mat",)),
    operation="scale",
  )

  mat_id = mujoco.mj_name2id(
    env.sim.mj_model, mujoco.mjtObj.mjOBJ_MATERIAL, "robot/test_mat"
  )
  expected = env.sim.get_default_field("mat_texrepeat")[mat_id] * 2.0
  texrepeat = env.sim.model.mat_texrepeat[:, mat_id, :]
  assert torch.allclose(texrepeat, expected.expand_as(texrepeat))


def test_mat_rgba_invalid_name(mat_env):
  """ValueError for unknown material name."""
  env = mat_env

  with pytest.raises(ValueError, match="nonexistent_material"):
    cfg = SceneEntityCfg("robot", material_names=("nonexistent_material",))
    cfg.resolve(env.scene)
    dr.mat_rgba(
      env,
      env_ids=None,
      ranges=(0.2, 0.8),
      asset_cfg=cfg,
    )


# geom_matid DR tests.

GEOM_MATID_XML = """
<mujoco>
  <asset>
    <material name="mat_a" rgba="1 0 0 1"/>
    <material name="mat_b" rgba="0 1 0 1"/>
    <material name="mat_c" rgba="0 0 1 1"/>
  </asset>
  <worldbody>
    <body name="base" pos="0 0 1">
      <freejoint name="free_joint"/>
      <geom name="geom1" type="box" size="0.1 0.1 0.1" mass="1.0" material="mat_a"/>
      <geom name="geom2" type="sphere" size="0.05" mass="0.3" material="mat_b"/>
    </body>
  </worldbody>
</mujoco>
"""


def _make_matid_env(device, num_envs=NUM_ENVS):
  entity_cfg = EntityCfg(spec_fn=lambda: mujoco.MjSpec.from_string(GEOM_MATID_XML))
  scene_cfg = SceneCfg(num_envs=num_envs, entities={"robot": entity_cfg})
  scene = Scene(scene_cfg, device)
  model = scene.compile()
  sim = Simulation(num_envs=num_envs, cfg=SimulationCfg(), model=model, device=device)
  scene.initialize(model, sim.model, sim.data)
  sim.expand_model_fields(("geom_matid",))
  return Env(scene, sim, device)


@pytest.fixture(scope="module")
def matid_env(device):
  return _make_matid_env(device)


def test_geom_matid_draws_from_pool(matid_env):
  """Assigned matids come from the pool and actually change from the defaults."""
  torch.manual_seed(42)
  env = matid_env
  robot = env.scene["robot"]
  asset_cfg = SceneEntityCfg("robot", geom_names=(".*",), material_names=(".*",))
  asset_cfg.resolve(env.scene)

  geom_ids = robot.indexing.geom_ids[asset_cfg.geom_ids]
  mat_ids = robot.indexing.mat_ids[asset_cfg.material_ids]
  original = env.sim.model.geom_matid[:, geom_ids].clone()

  dr.geom_matid(env, env_ids=None, asset_cfg=asset_cfg)

  assigned = env.sim.model.geom_matid[:, geom_ids]
  # Every assignment is a valid member of the selected pool.
  assert torch.all(torch.isin(assigned, mat_ids))
  # Randomization is not a no-op: values differ from the baked defaults, and at
  # least one geom draws a material it did not start with (proving the full pool
  # is sampled, not just the compile-time defaults).
  assert not torch.equal(assigned, original)
  assert torch.isin(assigned, torch.unique(original), invert=True).any()


def test_geom_matid_partial_env_ids(matid_env):
  """Randomizing subset of envs leaves others unchanged."""
  torch.manual_seed(42)
  env = matid_env
  robot = env.scene["robot"]
  asset_cfg = SceneEntityCfg("robot", geom_names=(".*",), material_names=(".*",))
  asset_cfg.resolve(env.scene)
  geom_ids = robot.indexing.geom_ids[asset_cfg.geom_ids]

  original = env.sim.model.geom_matid[:, geom_ids].clone()

  dr.geom_matid(
    env, env_ids=torch.tensor([0, 2], device=env.device), asset_cfg=asset_cfg
  )

  result = env.sim.model.geom_matid[:, geom_ids]
  assert torch.all(result[1] == original[1])
  assert torch.all(result[3] == original[3])


def test_geom_matid_invalid_material_name(matid_env):
  """Unknown material name is rejected during config resolution."""
  env = matid_env

  with pytest.raises(ValueError, match="nonexistent_material"):
    cfg = SceneEntityCfg(
      "robot", geom_names=(".*",), material_names=("nonexistent_material",)
    )
    cfg.resolve(env.scene)
    dr.geom_matid(env, env_ids=None, asset_cfg=cfg)


def test_geom_matid_empty_material_selection(matid_env):
  """geom_matid raises when the resolved material pool is empty."""
  env = matid_env
  cfg = SceneEntityCfg("robot", geom_names=(".*",), material_names=())
  cfg.resolve(env.scene)

  with pytest.raises(ValueError, match="No materials selected"):
    dr.geom_matid(env, env_ids=None, asset_cfg=cfg)


def test_geom_matid_shared_random(matid_env):
  """All geoms within same env get same material, envs differ."""
  torch.manual_seed(42)
  env = matid_env
  robot = env.scene["robot"]
  asset_cfg = SceneEntityCfg("robot", geom_names=(".*",), material_names=(".*",))
  asset_cfg.resolve(env.scene)
  geom_ids = robot.indexing.geom_ids[asset_cfg.geom_ids]

  dr.geom_matid(env, env_ids=None, asset_cfg=asset_cfg, shared_random=True)

  assigned = env.sim.model.geom_matid[:, geom_ids]
  for env_idx in range(env.num_envs):
    env_matids = assigned[env_idx]
    assert torch.all(env_matids == env_matids[0])

  assert len(torch.unique(assigned[:, 0])) > 1


# mat_texid DR tests.

MAT_TEXID_XML = """
<mujoco>
  <asset>
    <texture name="tex_a" type="2d" builtin="flat" rgb1="1 0 0" width="4" height="4"/>
    <texture name="tex_b" type="2d" builtin="flat" rgb1="0 1 0" width="4" height="4"/>
    <texture name="tex_c" type="2d" builtin="flat" rgb1="0 0 1" width="4" height="4"/>
    <texture name="tex_d" type="2d" builtin="flat" rgb1="1 1 0" width="4" height="4"/>
    <material name="mat_a" texture="tex_a"/>
    <material name="mat_b" texture="tex_b"/>
  </asset>
  <worldbody>
    <body name="base" pos="0 0 1">
      <freejoint name="free_joint"/>
      <geom name="geom1" type="box" size="0.1 0.1 0.1" mass="1.0" material="mat_a"/>
      <geom name="geom2" type="sphere" size="0.05" mass="0.3" material="mat_b"/>
    </body>
  </worldbody>
</mujoco>
"""


def _make_texid_env(device, num_envs=NUM_ENVS):
  entity_cfg = EntityCfg(spec_fn=lambda: mujoco.MjSpec.from_string(MAT_TEXID_XML))
  scene_cfg = SceneCfg(num_envs=num_envs, entities={"robot": entity_cfg})
  scene = Scene(scene_cfg, device)
  model = scene.compile()
  sim = Simulation(num_envs=num_envs, cfg=SimulationCfg(), model=model, device=device)
  scene.initialize(model, sim.model, sim.data)
  sim.expand_model_fields(("mat_texid",))
  return Env(scene, sim, device)


@pytest.fixture
def texid_env(device):
  # Function-scoped so tests compare against pristine mat_texid defaults.
  return _make_texid_env(device)


def test_mat_texid_draws_from_pool(texid_env):
  """Assigned texids come from the pool and actually change from the defaults."""
  torch.manual_seed(42)
  env = texid_env
  robot = env.scene["robot"]
  asset_cfg = SceneEntityCfg("robot", material_names=(".*",), texture_names=(".*",))
  asset_cfg.resolve(env.scene)

  mat_ids = robot.indexing.mat_ids[asset_cfg.material_ids]
  tex_ids = robot.indexing.tex_ids[asset_cfg.texture_ids]
  role = mujoco.mjtTextureRole.mjTEXROLE_RGB.value
  original = env.sim.model.mat_texid[:, mat_ids, role].clone()

  dr.mat_texid(env, env_ids=None, asset_cfg=asset_cfg)

  assigned = env.sim.model.mat_texid[:, mat_ids, role]
  assert torch.all(torch.isin(assigned, tex_ids))
  assert not torch.equal(assigned, original)
  assert torch.isin(assigned, torch.unique(original), invert=True).any()


def test_mat_texid_respects_texture_subset(texid_env):
  """Textures outside the selection are never sampled."""
  torch.manual_seed(42)
  env = texid_env
  robot = env.scene["robot"]
  # The pool excludes tex_a and tex_b, which are the two materials' defaults.
  asset_cfg = SceneEntityCfg(
    "robot", material_names=(".*",), texture_names=("tex_c", "tex_d")
  )
  asset_cfg.resolve(env.scene)

  mat_ids = robot.indexing.mat_ids[asset_cfg.material_ids]
  tex_ids = robot.indexing.tex_ids[asset_cfg.texture_ids]
  role = mujoco.mjtTextureRole.mjTEXROLE_RGB.value
  excluded = torch.unique(env.sim.model.mat_texid[:, mat_ids, role])

  dr.mat_texid(env, env_ids=None, asset_cfg=asset_cfg)

  assigned = env.sim.model.mat_texid[:, mat_ids, role]
  assert torch.all(torch.isin(assigned, tex_ids))
  assert not torch.isin(assigned, excluded).any()
  assert len(torch.unique(assigned)) >= 2


def test_mat_texid_partial_env_ids(texid_env):
  """Randomizing subset of envs leaves others unchanged."""
  torch.manual_seed(42)
  env = texid_env
  robot = env.scene["robot"]
  asset_cfg = SceneEntityCfg(
    "robot", material_names=(".*",), texture_names=("tex_c", "tex_d")
  )
  asset_cfg.resolve(env.scene)
  mat_ids = robot.indexing.mat_ids[asset_cfg.material_ids]
  tex_ids = robot.indexing.tex_ids[asset_cfg.texture_ids]
  role = mujoco.mjtTextureRole.mjTEXROLE_RGB.value

  original = env.sim.model.mat_texid[:, mat_ids, role].clone()

  dr.mat_texid(
    env, env_ids=torch.tensor([0, 2], device=env.device), asset_cfg=asset_cfg
  )

  result = env.sim.model.mat_texid[:, mat_ids, role]
  # Envs 0, 2 changed, the pool excludes their defaults.
  assert torch.all(torch.isin(result[0], tex_ids))
  assert torch.all(torch.isin(result[2], tex_ids))
  assert not torch.equal(result[0], original[0])
  assert not torch.equal(result[2], original[2])
  # Envs 1, 3 unchanged.
  assert torch.all(result[1] == original[1])
  assert torch.all(result[3] == original[3])


def test_mat_texid_invalid_texture_name(texid_env):
  """Unknown texture name is rejected during config resolution."""
  env = texid_env

  with pytest.raises(ValueError, match="nonexistent_texture"):
    cfg = SceneEntityCfg(
      "robot", material_names=(".*",), texture_names=("nonexistent_texture",)
    )
    cfg.resolve(env.scene)
    dr.mat_texid(env, env_ids=None, asset_cfg=cfg)


def test_mat_texid_empty_texture_selection(texid_env):
  """mat_texid raises when the resolved texture pool is empty."""
  env = texid_env
  cfg = SceneEntityCfg("robot", material_names=(".*",), texture_names=())
  cfg.resolve(env.scene)

  with pytest.raises(ValueError, match="No textures selected"):
    dr.mat_texid(env, env_ids=None, asset_cfg=cfg)


def test_mat_texid_shared_random(texid_env):
  """All materials within same env get same texture, envs differ."""
  torch.manual_seed(42)
  env = texid_env
  robot = env.scene["robot"]
  asset_cfg = SceneEntityCfg("robot", material_names=(".*",), texture_names=(".*",))
  asset_cfg.resolve(env.scene)
  mat_ids = robot.indexing.mat_ids[asset_cfg.material_ids]
  role = mujoco.mjtTextureRole.mjTEXROLE_RGB.value

  dr.mat_texid(env, env_ids=None, asset_cfg=asset_cfg, shared_random=True)

  assigned = env.sim.model.mat_texid[:, mat_ids, role]
  for env_idx in range(env.num_envs):
    env_texids = assigned[env_idx]
    assert torch.all(env_texids == env_texids[0])

  assert len(torch.unique(assigned[:, 0])) > 1


# pair_friction tests.

PAIR_XML = """
<mujoco>
  <worldbody>
    <body name="base" pos="0 0 1">
      <freejoint name="free_joint"/>
      <geom name="base_geom" type="box" size="0.1 0.1 0.1" mass="1.0"/>
      <body name="foot" pos="0 0 -0.15">
        <geom name="foot_geom" type="box" size="0.05 0.05 0.02" mass="0.1"/>
      </body>
    </body>
    <geom name="floor" type="plane" size="2 2 0.1"/>
  </worldbody>
  <contact>
    <pair name="base_floor" geom1="base_geom" geom2="floor"
          friction="0.5 0.5 0.005 0.0001 0.0001"/>
    <pair name="foot_floor" geom1="foot_geom" geom2="floor"
          friction="0.3 0.3 0.003 0.0002 0.0002"/>
  </contact>
</mujoco>
"""


@pytest.fixture(scope="module")
def pair_env(device):
  entity_cfg = EntityCfg(spec_fn=lambda: mujoco.MjSpec.from_string(PAIR_XML))
  scene_cfg = SceneCfg(num_envs=NUM_ENVS, entities={"robot": entity_cfg})
  scene = Scene(scene_cfg, device)
  model = scene.compile()
  sim = Simulation(num_envs=NUM_ENVS, cfg=SimulationCfg(), model=model, device=device)
  scene.initialize(model, sim.model, sim.data)
  sim.expand_model_fields(("pair_friction",))
  return Env(scene, sim, device)


def test_pair_friction_abs(pair_env):
  """Values set directly within range, diverse across envs."""
  torch.manual_seed(42)
  env = pair_env
  robot = env.scene["robot"]

  dr.pair_friction(env, env_ids=None, ranges=(0.4, 1.5), operation="abs")

  friction = env.sim.model.pair_friction[:, robot.indexing.pair_ids, 0]
  assert torch.all((friction >= 0.4) & (friction <= 1.5))
  assert len(torch.unique(friction)) >= 2


def test_pair_friction_scale(pair_env):
  """Values = default * sample; no accumulation after 3 calls."""
  env = pair_env
  robot = env.scene["robot"]
  pair_idx = robot.indexing.pair_ids[0]
  default_val = env.sim.get_default_field("pair_friction")[pair_idx, 0].item()

  for _ in range(3):
    dr.pair_friction(
      env,
      env_ids=None,
      ranges=(2.0, 2.0),
      operation="scale",
      asset_cfg=SceneEntityCfg("robot", pair_ids=[0]),
      axes=[0],
    )

  result = env.sim.model.pair_friction[:, pair_idx, 0]
  assert torch.allclose(result, torch.full_like(result, default_val * 2.0), atol=1e-5)


def test_pair_friction_add(pair_env):
  """Values = default + sample; no accumulation after 3 calls."""
  env = pair_env
  robot = env.scene["robot"]
  pair_idx = robot.indexing.pair_ids[0]
  default_val = env.sim.get_default_field("pair_friction")[pair_idx, 0].item()

  for _ in range(3):
    dr.pair_friction(
      env,
      env_ids=None,
      ranges=(0.1, 0.1),
      operation="add",
      asset_cfg=SceneEntityCfg("robot", pair_ids=[0]),
      axes=[0],
    )

  result = env.sim.model.pair_friction[:, pair_idx, 0]
  assert torch.allclose(result, torch.full_like(result, default_val + 0.1), atol=1e-5)


def test_pair_friction_axes_selectivity(pair_env):
  """Only the targeted axis changes; others stay at their original values."""
  torch.manual_seed(7)
  env = pair_env
  robot = env.scene["robot"]

  before = env.sim.model.pair_friction.clone()

  dr.pair_friction(env, env_ids=None, ranges=(0.8, 1.2), operation="abs", axes=[0])

  after = env.sim.model.pair_friction
  pair_ids = robot.indexing.pair_ids
  # Axis 0 changed.
  assert not torch.allclose(after[:, pair_ids, 0], before[:, pair_ids, 0])
  # Axes 1-4 unchanged.
  for ax in range(1, 5):
    assert torch.allclose(after[:, pair_ids, ax], before[:, pair_ids, ax])


def test_pair_friction_multi_axis(pair_env):
  """axes=[0, 1] randomizes both tangent components; axes 2-4 stay unchanged."""
  torch.manual_seed(17)
  env = pair_env
  robot = env.scene["robot"]

  before = env.sim.model.pair_friction.clone()

  dr.pair_friction(env, env_ids=None, ranges=(0.4, 1.0), operation="abs", axes=[0, 1])

  after = env.sim.model.pair_friction
  pair_ids = robot.indexing.pair_ids
  # Both tangent axes changed.
  assert not torch.allclose(after[:, pair_ids, 0], before[:, pair_ids, 0])
  assert not torch.allclose(after[:, pair_ids, 1], before[:, pair_ids, 1])
  # Axes 2-4 unchanged.
  for ax in range(2, 5):
    assert torch.allclose(after[:, pair_ids, ax], before[:, pair_ids, ax])
  # Axis 0 and axis 1 are sampled independently (not forced equal).
  assert not torch.allclose(after[:, pair_ids, 0], after[:, pair_ids, 1])


def test_pair_friction_partial_env_ids(pair_env):
  """Randomizing a subset of envs leaves the others unchanged."""
  torch.manual_seed(99)
  env = pair_env
  robot = env.scene["robot"]

  before = env.sim.model.pair_friction.clone()
  randomized_ids = torch.tensor([0, 1], device=env.device)

  dr.pair_friction(
    env, env_ids=randomized_ids, ranges=(2.0, 3.0), operation="abs", axes=[0]
  )

  after = env.sim.model.pair_friction
  pair_ids = robot.indexing.pair_ids
  unchanged_ids = torch.tensor([2, 3], device=env.device)
  assert torch.allclose(
    after[unchanged_ids[:, None], pair_ids[None, :], :],
    before[unchanged_ids[:, None], pair_ids[None, :], :],
  )


def test_pair_friction_shared_random(pair_env):
  """All pairs in the same env get the same value; envs differ."""
  torch.manual_seed(5)
  env = pair_env
  robot = env.scene["robot"]

  dr.pair_friction(
    env, env_ids=None, ranges=(0.5, 2.0), operation="abs", axes=[0], shared_random=True
  )

  pair_ids = robot.indexing.pair_ids
  friction = env.sim.model.pair_friction[:, pair_ids, 0]
  # Within each env, all pairs have the same value.
  for e in range(env.num_envs):
    row = friction[e]
    assert torch.allclose(row, row[0].expand_as(row), atol=1e-6)
  # Across envs, values differ.
  assert not torch.allclose(friction[0], friction[1])


def test_pair_friction_by_name(pair_env):
  """Targeting pair by name only modifies that pair."""
  torch.manual_seed(11)
  env = pair_env
  robot = env.scene["robot"]

  foot_cfg = SceneEntityCfg("robot", pair_names=("foot_floor",))
  foot_cfg.resolve(env.scene)
  base_cfg = SceneEntityCfg("robot", pair_names=("base_floor",))
  base_cfg.resolve(env.scene)
  foot_pair_id = robot.indexing.pair_ids[foot_cfg.pair_ids]
  base_pair_id = robot.indexing.pair_ids[base_cfg.pair_ids]

  before = env.sim.model.pair_friction.clone()
  dr.pair_friction(
    env, env_ids=None, ranges=(1.5, 2.5), operation="abs", asset_cfg=foot_cfg
  )

  after = env.sim.model.pair_friction
  # foot_floor changed.
  assert not torch.allclose(after[:, foot_pair_id, 0], before[:, foot_pair_id, 0])
  # base_floor unchanged.
  assert torch.allclose(after[:, base_pair_id, :], before[:, base_pair_id, :])


def test_pair_friction_string_ranges(pair_env):
  """Dict with pair name patterns randomizes the matched pairs correctly."""
  torch.manual_seed(13)
  env = pair_env
  robot = env.scene["robot"]

  dr.pair_friction(
    env,
    env_ids=None,
    ranges={"base_floor": (1.0, 1.0), "foot_floor": (0.5, 0.5)},
    operation="abs",
    axes=[0],
  )

  base_cfg = SceneEntityCfg("robot", pair_names=("base_floor",))
  base_cfg.resolve(env.scene)
  foot_cfg = SceneEntityCfg("robot", pair_names=("foot_floor",))
  foot_cfg.resolve(env.scene)
  base_id = robot.indexing.pair_ids[base_cfg.pair_ids]
  foot_id = robot.indexing.pair_ids[foot_cfg.pair_ids]
  assert torch.allclose(
    env.sim.model.pair_friction[:, base_id, 0],
    torch.ones(env.num_envs, device=env.device),
    atol=1e-5,
  )
  assert torch.allclose(
    env.sim.model.pair_friction[:, foot_id, 0],
    torch.full((env.num_envs,), 0.5, device=env.device),
    atol=1e-5,
  )


def test_pair_friction_isotropic_tangent(pair_env):
  """isotropic=True mirrors tangent2 = tangent1; spin and roll axes unchanged."""
  torch.manual_seed(21)
  env = pair_env
  robot = env.scene["robot"]
  pair_ids = robot.indexing.pair_ids

  before = env.sim.model.pair_friction.clone()
  dr.pair_friction(
    env, env_ids=None, ranges=(0.4, 1.0), operation="abs", axes=[0], isotropic=True
  )

  after = env.sim.model.pair_friction
  # Axis 0 was randomized into range.
  assert torch.all((after[:, pair_ids, 0] >= 0.4) & (after[:, pair_ids, 0] <= 1.0))
  # Axis 1 equals axis 0 exactly.
  assert torch.allclose(after[:, pair_ids, 1], after[:, pair_ids, 0])
  # Axes 2, 3, 4 untouched.
  for ax in (2, 3, 4):
    assert torch.allclose(after[:, pair_ids, ax], before[:, pair_ids, ax])


def test_pair_friction_isotropic_shared(pair_env):
  """isotropic=True + shared_random=True: all pairs same value, tangent2 == tangent1."""
  torch.manual_seed(33)
  env = pair_env
  robot = env.scene["robot"]
  pair_ids = robot.indexing.pair_ids

  dr.pair_friction(
    env,
    env_ids=None,
    ranges=(0.4, 1.0),
    operation="abs",
    axes=[0],
    shared_random=True,
    isotropic=True,
  )

  friction = env.sim.model.pair_friction[:, pair_ids, :]
  # All pairs in each env share the same tangent1 value.
  for e in range(env.num_envs):
    assert torch.allclose(
      friction[e, :, 0], friction[e, 0, 0].expand(len(pair_ids)), atol=1e-6
    )
  # Across envs, values differ.
  assert not torch.allclose(friction[0, 0, 0], friction[1, 0, 0])
  # tangent2 == tangent1 everywhere.
  assert torch.allclose(friction[:, :, 1], friction[:, :, 0])


def test_pair_friction_isotropic_roll(pair_env):
  """isotropic=True with roll axes mirrors roll2 = roll1; tangent axes unchanged."""
  torch.manual_seed(41)
  env = pair_env
  robot = env.scene["robot"]
  pair_ids = robot.indexing.pair_ids

  before = env.sim.model.pair_friction.clone()
  dr.pair_friction(
    env, env_ids=None, ranges=(0.001, 0.01), operation="abs", axes=[3], isotropic=True
  )

  after = env.sim.model.pair_friction
  # Axis 3 was randomized.
  assert not torch.allclose(after[:, pair_ids, 3], before[:, pair_ids, 3])
  # Axis 4 equals axis 3 exactly.
  assert torch.allclose(after[:, pair_ids, 4], after[:, pair_ids, 3])
  # Tangent and spin axes untouched.
  for ax in (0, 1, 2):
    assert torch.allclose(after[:, pair_ids, ax], before[:, pair_ids, ax])


def test_pair_friction_isotropic_roll_dict_ranges(pair_env):
  """isotropic=True with dict-int ranges containing roll axis triggers roll2 = roll1."""
  torch.manual_seed(53)
  env = pair_env
  robot = env.scene["robot"]
  pair_ids = robot.indexing.pair_ids

  before = env.sim.model.pair_friction.clone()
  dr.pair_friction(
    env,
    env_ids=None,
    ranges={3: (0.001, 0.01)},
    operation="abs",
    isotropic=True,
  )

  after = env.sim.model.pair_friction
  # Axis 3 was randomized.
  assert not torch.allclose(after[:, pair_ids, 3], before[:, pair_ids, 3])
  # Axis 4 equals axis 3 exactly (dict-int ranges must trigger the roll copy).
  assert torch.allclose(after[:, pair_ids, 4], after[:, pair_ids, 3])
  # Tangent and spin axes untouched.
  for ax in (0, 1, 2):
    assert torch.allclose(after[:, pair_ids, ax], before[:, pair_ids, ax])


def test_pair_friction_invalid_name(pair_env):
  """ValueError for unknown pair name."""
  env = pair_env

  with pytest.raises(ValueError, match="nonexistent_pair"):
    cfg = SceneEntityCfg("robot", pair_names=("nonexistent_pair",))
    cfg.resolve(env.scene)
