"""Tests for EventManager, special dr.* functions, and recomputation."""

from unittest.mock import Mock

import mujoco
import pytest
import torch
from conftest import get_test_device

from mjlab import actuator
from mjlab.entity import EntityCfg
from mjlab.envs.mdp import dr, events
from mjlab.managers.event_manager import (
  EventManager,
  EventTermCfg,
  RecomputeLevel,
)
from mjlab.managers.scene_entity_config import SceneEntityCfg
from mjlab.scene import Scene, SceneCfg
from mjlab.sim.sim import Simulation, SimulationCfg

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="5.0"/>
      <body name="link1" pos="0.2 0 0">
        <joint name="joint1" type="hinge" axis="0 0 1" range="0 1.57"
          armature="0.1"/>
        <geom name="link1_geom" type="box" size="0.05 0.05 0.2" mass="1.0"/>
      </body>
      <body name="link2" pos="-0.2 0 0">
        <joint name="joint2" type="hinge" axis="0 0 1" range="0 1.57"
          armature="0.1"/>
        <geom name="link2_geom" type="box" size="0.05 0.05 0.2" mass="1.0"/>
      </body>
    </body>
  </worldbody>
</mujoco>
"""

NUM_ENVS = 2


@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_env(device, fields, num_envs=NUM_ENVS):
  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)
  sim.expand_model_fields(fields)

  return Env(scene, sim, device)


# ===========================================================================
# Section 1: EventManager
# ===========================================================================


def test_dr_fields_registered_in_event_manager(device):
  """@requires_model_fields functions populate manager.domain_randomization_fields."""
  env = Mock()
  env.num_envs = 4
  env.device = device
  env.scene = {}
  env.sim = Mock()

  cfg = {
    "friction_dr": EventTermCfg(
      mode="startup",
      func=dr.geom_friction,
      params={"ranges": (0.3, 1.2)},
    ),
    "damping_dr": EventTermCfg(
      mode="reset",
      func=dr.joint_damping,
      params={"ranges": (0.1, 0.5)},
    ),
    "pd_gains": EventTermCfg(
      mode="reset",
      func=dr.pd_gains,
      params={"kp_range": (0.8, 1.2), "kd_range": (0.8, 1.2)},
    ),
    "effort_limits": EventTermCfg(
      mode="reset",
      func=dr.effort_limits,
      params={"effort_limit_range": (0.8, 1.2)},
    ),
    "regular_event": EventTermCfg(
      mode="reset",
      func=events.reset_joints_by_offset,
      params={"position_range": (-0.1, 0.1), "velocity_range": (0.0, 0.0)},
    ),
  }

  manager = EventManager(cfg, env)

  assert "geom_friction" in manager.domain_randomization_fields
  assert "dof_damping" in manager.domain_randomization_fields
  assert "actuator_gainprm" in manager.domain_randomization_fields
  assert "actuator_biasprm" in manager.domain_randomization_fields
  assert "actuator_forcerange" in manager.domain_randomization_fields
  assert "jnt_actfrcrange" in manager.domain_randomization_fields
  assert "tendon_actfrcrange" in manager.domain_randomization_fields
  assert len(manager.domain_randomization_fields) == 7


def test_recompute_fields_registered_in_event_manager(device):
  """Recomputed model fields are expanded per environment."""
  env = Mock()
  env.num_envs = 4
  env.device = device
  env.scene = {}
  env.sim = Mock()

  cfg = {
    "body_com": EventTermCfg(
      mode="startup",
      func=dr.body_com_offset,
      params={"ranges": (-0.01, 0.01)},
    ),
  }

  manager = EventManager(cfg, env)

  assert manager.domain_randomization_fields == (
    "body_ipos",
    "body_subtreemass",
    "dof_invweight0",
    "body_invweight0",
    "tendon_length0",
    "tendon_invweight0",
    "actuator_acc0",
  )


def test_recompute_level_ordering():
  """IntEnum max() gives strongest level."""
  L = RecomputeLevel
  assert max(L.none, L.set_const_0) == L.set_const_0
  assert max(L.set_const_0, L.set_const) == L.set_const
  assert max(L.set_const, L.set_const_fixed) == L.set_const
  assert max(L.none, L.none) == L.none


def test_recompute_not_called_when_no_dr_fired(device):
  """Non-DR events don't trigger recompute."""
  env = Mock()
  env.num_envs = 2
  env.device = device
  env.scene = {}
  env.sim = Mock()
  env.sim.recompute_constants = Mock()

  cfg = {
    "regular_event": EventTermCfg(
      mode="reset",
      func=lambda env, env_ids: None,
      params={},
    ),
  }

  manager = EventManager(cfg, env)
  manager.apply(
    "reset",
    env_ids=torch.tensor([0], device=device),
    global_env_step_count=1,
  )

  env.sim.recompute_constants.assert_not_called()


# ===========================================================================
# Section 2: PD gains & effort limits
# ===========================================================================


def _make_pd_env(device, num_envs=2):
  """Create a mock env with Builtin, Xml, and Ideal actuators for PD tests."""
  env = Mock()
  env.num_envs = num_envs
  env.device = device

  mock_entity = Mock()

  builtin = Mock(spec=actuator.BuiltinPositionActuator)
  builtin.ctrl_ids = torch.tensor([0, 1], device=device)
  builtin.global_ctrl_ids = torch.tensor([0, 1], device=device)

  xml = Mock(spec=actuator.XmlActuator)
  xml.command_field = "position"
  xml.ctrl_ids = torch.tensor([2, 3], device=device)
  xml.global_ctrl_ids = torch.tensor([2, 3], device=device)

  ideal = Mock(spec=actuator.IdealPdActuator)
  ideal.ctrl_ids = torch.tensor([4, 5], device=device)
  ideal.global_ctrl_ids = torch.tensor([4, 5], device=device)
  ideal.stiffness = torch.tensor([[100.0, 100.0]] * num_envs, device=device)
  ideal.damping = torch.tensor([[10.0, 10.0]] * num_envs, device=device)
  ideal.default_stiffness = torch.tensor([[100.0, 100.0]] * num_envs, device=device)
  ideal.default_damping = torch.tensor([[10.0, 10.0]] * num_envs, device=device)
  ideal.set_gains = actuator.IdealPdActuator.set_gains.__get__(ideal)

  mock_entity.actuators = [builtin, xml, ideal]
  env.scene = {"robot": mock_entity}

  env.sim = Mock()
  env.sim.model = Mock()
  env.sim.model.actuator_gainprm = torch.ones((num_envs, 6, 10), device=device) * 50.0
  env.sim.model.actuator_biasprm = torch.zeros((num_envs, 6, 10), device=device)
  env.sim.model.actuator_biasprm[:, :, 1] = -50.0
  env.sim.model.actuator_biasprm[:, :, 2] = -5.0

  default_gainprm = torch.ones((6, 10), device=device) * 50.0
  default_biasprm = torch.zeros((6, 10), device=device)
  default_biasprm[:, 1] = -50.0
  default_biasprm[:, 2] = -5.0
  defaults = {
    "actuator_gainprm": default_gainprm,
    "actuator_biasprm": default_biasprm,
  }
  env.sim.get_default_field = lambda f: defaults[f]

  return env, ideal


def _make_effort_env(device, num_envs=2):
  """Create a mock env with Builtin, Xml, and Ideal actuators for effort tests."""
  env = Mock()
  env.num_envs = num_envs
  env.device = device

  mock_entity = Mock()

  builtin = Mock(spec=actuator.BuiltinPositionActuator)
  builtin.ctrl_ids = torch.tensor([0, 1], device=device)
  builtin.global_ctrl_ids = torch.tensor([0, 1], device=device)

  xml = Mock(spec=actuator.XmlActuator)
  xml.command_field = "position"
  xml.ctrl_ids = torch.tensor([2, 3], device=device)
  xml.global_ctrl_ids = torch.tensor([2, 3], device=device)

  ideal = Mock(spec=actuator.IdealPdActuator)
  ideal.ctrl_ids = torch.tensor([4, 5], device=device)
  ideal.global_ctrl_ids = torch.tensor([4, 5], device=device)
  ideal.force_limit = torch.tensor([[50.0, 50.0]] * num_envs, device=device)
  ideal.default_force_limit = torch.tensor([[50.0, 50.0]] * num_envs, device=device)
  ideal.set_effort_limit = actuator.IdealPdActuator.set_effort_limit.__get__(ideal)

  mock_entity.actuators = [builtin, xml, ideal]
  env.scene = {"robot": mock_entity}

  env.sim = Mock()
  env.sim.model = Mock()
  env.sim.model.actuator_forcerange = torch.zeros((num_envs, 6, 2), device=device)
  env.sim.model.actuator_forcerange[:, :, 0] = -100.0
  env.sim.model.actuator_forcerange[:, :, 1] = 100.0

  default_forcerange = torch.zeros((6, 2), device=device)
  default_forcerange[:, 0] = -100.0
  default_forcerange[:, 1] = 100.0
  env.sim.get_default_field = lambda field: default_forcerange

  return env, ideal


@pytest.mark.parametrize(
  "operation, env_id, kp, kd, expected_gainprm, expected_biasprm_1,"
  " expected_biasprm_2, expected_stiffness, expected_damping",
  [
    ("scale", 0, 1.5, 2.0, 75.0, -75.0, -10.0, 150.0, 20.0),
    ("abs", 1, 200.0, 25.0, 200.0, -200.0, -25.0, 200.0, 25.0),
  ],
)
def test_pd_gains(
  device,
  operation,
  env_id,
  kp,
  kd,
  expected_gainprm,
  expected_biasprm_1,
  expected_biasprm_2,
  expected_stiffness,
  expected_damping,
):
  """PD gains scale/abs on Builtin, Xml, and Ideal actuators."""
  env, ideal = _make_pd_env(device)

  dr.pd_gains(
    env,
    torch.tensor([env_id], device=device),
    kp_range=(kp, kp),
    kd_range=(kd, kd),
    asset_cfg=SceneEntityCfg("robot"),
    operation=operation,
  )

  d = device
  # Builtin actuator.
  assert torch.allclose(
    env.sim.model.actuator_gainprm[env_id, [0, 1], 0],
    torch.tensor([expected_gainprm] * 2, device=d),
  )
  assert torch.allclose(
    env.sim.model.actuator_biasprm[env_id, [0, 1], 1],
    torch.tensor([expected_biasprm_1] * 2, device=d),
  )
  assert torch.allclose(
    env.sim.model.actuator_biasprm[env_id, [0, 1], 2],
    torch.tensor([expected_biasprm_2] * 2, device=d),
  )
  # Xml actuator.
  assert torch.allclose(
    env.sim.model.actuator_gainprm[env_id, [2, 3], 0],
    torch.tensor([expected_gainprm] * 2, device=d),
  )
  # Ideal actuator.
  assert torch.allclose(
    ideal.stiffness[env_id],
    torch.tensor([expected_stiffness] * 2, device=d),
  )
  assert torch.allclose(
    ideal.damping[env_id],
    torch.tensor([expected_damping] * 2, device=d),
  )


def test_pd_gains_multi_env(device):
  """Independent values per environment."""
  env, _ = _make_pd_env(device)

  torch.manual_seed(42)
  dr.pd_gains(
    env,
    torch.tensor([0, 1], device=device),
    kp_range=(0.5, 2.0),
    kd_range=(0.5, 2.0),
    asset_cfg=SceneEntityCfg("robot"),
    operation="scale",
  )

  gains = env.sim.model.actuator_gainprm[:, :2, 0]
  assert (gains != 50.0).all()
  assert not torch.allclose(gains[0], gains[1])


@pytest.mark.parametrize(
  "operation, env_id, limit, expected_lower, expected_upper, expected_ideal",
  [
    ("scale", 0, 2.0, -200.0, 200.0, 100.0),
    ("abs", 1, 150.0, -150.0, 150.0, 150.0),
  ],
)
def test_effort_limits(
  device, operation, env_id, limit, expected_lower, expected_upper, expected_ideal
):
  """Effort limits scale/abs on Builtin, Xml, and Ideal actuators."""
  env, ideal = _make_effort_env(device)

  dr.effort_limits(
    env,
    torch.tensor([env_id], device=device),
    effort_limit_range=(limit, limit),
    asset_cfg=SceneEntityCfg("robot"),
    operation=operation,
  )

  d = device
  # Builtin.
  assert torch.allclose(
    env.sim.model.actuator_forcerange[env_id, [0, 1], 0],
    torch.tensor([expected_lower] * 2, device=d),
  )
  assert torch.allclose(
    env.sim.model.actuator_forcerange[env_id, [0, 1], 1],
    torch.tensor([expected_upper] * 2, device=d),
  )
  # Xml.
  assert torch.allclose(
    env.sim.model.actuator_forcerange[env_id, [2, 3], 0],
    torch.tensor([expected_lower] * 2, device=d),
  )
  # Ideal.
  assert torch.allclose(
    ideal.force_limit[env_id],
    torch.tensor([expected_ideal] * 2, device=d),
  )


def test_effort_limits_multi_env(device):
  """Independent values per environment."""
  env, _ = _make_effort_env(device)

  torch.manual_seed(42)
  dr.effort_limits(
    env,
    torch.tensor([0, 1], device=device),
    effort_limit_range=(0.5, 2.0),
    asset_cfg=SceneEntityCfg("robot"),
    operation="scale",
  )

  upper = env.sim.model.actuator_forcerange[:, :2, 1]
  assert (upper != 100.0).all()
  assert not torch.allclose(upper[0], upper[1])


def test_effort_limits_scale_no_accumulation(device):
  """Scale operation uses defaults (no accumulation)."""
  env, _ = _make_effort_env(device)

  for _ in range(3):
    dr.effort_limits(
      env,
      env_ids=None,
      effort_limit_range=(2.0, 2.0),
      asset_cfg=SceneEntityCfg("robot"),
      operation="scale",
    )

  actual_upper = env.sim.model.actuator_forcerange[0, 0, 1].item()
  assert abs(actual_upper - 200.0) < 1e-5


def test_pd_gains_accepts_operation_object(device):
  """dr.scale / dr.abs Operation objects produce the same result as strings."""
  env_str, ideal_str = _make_pd_env(device)
  env_obj, ideal_obj = _make_pd_env(device)

  ids = torch.tensor([0], device=device)
  kwargs = dict(
    kp_range=(1.5, 1.5), kd_range=(2.0, 2.0), asset_cfg=SceneEntityCfg("robot")
  )

  torch.manual_seed(0)
  dr.pd_gains(env_str, ids, operation="scale", **kwargs)
  torch.manual_seed(0)
  dr.pd_gains(env_obj, ids, operation=dr.scale, **kwargs)

  assert torch.allclose(
    env_str.sim.model.actuator_gainprm[0], env_obj.sim.model.actuator_gainprm[0]
  )
  assert torch.allclose(ideal_str.stiffness, ideal_obj.stiffness)


def test_effort_limits_accepts_operation_object(device):
  """dr.abs Operation object produces the same result as the string."""
  env_str, ideal_str = _make_effort_env(device)
  env_obj, ideal_obj = _make_effort_env(device)

  ids = torch.tensor([0], device=device)
  kwargs = dict(effort_limit_range=(150.0, 150.0), asset_cfg=SceneEntityCfg("robot"))

  dr.effort_limits(env_str, ids, operation="abs", **kwargs)
  dr.effort_limits(env_obj, ids, operation=dr.abs, **kwargs)

  assert torch.allclose(
    env_str.sim.model.actuator_forcerange[0], env_obj.sim.model.actuator_forcerange[0]
  )
  assert torch.allclose(ideal_str.force_limit, ideal_obj.force_limit)


def test_pd_gains_rejects_unsupported_operation(device):
  """Operations other than scale/abs raise ValueError."""
  env, _ = _make_pd_env(device)
  ids = torch.tensor([0], device=device)

  with pytest.raises(ValueError, match="only supports 'scale' and 'abs'"):
    dr.pd_gains(env, ids, kp_range=(1.0, 1.0), kd_range=(1.0, 1.0), operation=dr.add)


def test_effort_limits_rejects_unsupported_operation(device):
  """Operations other than scale/abs raise ValueError."""
  env, _ = _make_effort_env(device)
  ids = torch.tensor([0], device=device)

  with pytest.raises(ValueError, match="only supports 'scale' and 'abs'"):
    dr.effort_limits(env, ids, effort_limit_range=(1.0, 1.0), operation=dr.add)


# ===========================================================================
# Section 3: Other events
# ===========================================================================


def test_reset_joints_by_offset(device):
  """reset_joints_by_offset applies offsets and respects joint limits."""
  env = Mock()
  env.num_envs = 2
  env.device = device

  mock_entity = Mock()
  mock_entity.data.default_joint_pos = torch.zeros((2, 3), device=device)
  mock_entity.data.default_joint_vel = torch.zeros((2, 3), device=device)
  mock_entity.data.soft_joint_pos_limits = torch.tensor(
    [
      [[-0.5, 0.5], [-0.5, 0.5], [-0.5, 0.5]],
      [[-0.5, 0.5], [-0.5, 0.5], [-0.5, 0.5]],
    ],
    device=device,
  )
  mock_entity.write_joint_state_to_sim = Mock()
  env.scene = {"robot": mock_entity}

  # Normal offset.
  events.reset_joints_by_offset(
    env,
    torch.tensor([0], device=device),
    position_range=(0.3, 0.3),
    velocity_range=(0.2, 0.2),
    asset_cfg=SceneEntityCfg("robot", joint_ids=slice(None)),
  )

  call_args = mock_entity.write_joint_state_to_sim.call_args
  joint_pos, joint_vel = call_args[0][0], call_args[0][1]
  assert torch.allclose(joint_pos, torch.ones_like(joint_pos) * 0.3)
  assert torch.allclose(joint_vel, torch.ones_like(joint_vel) * 0.2)

  # Clamping when offset exceeds limits.
  events.reset_joints_by_offset(
    env,
    torch.tensor([1], device=device),
    position_range=(1.0, 1.0),
    velocity_range=(0.0, 0.0),
    asset_cfg=SceneEntityCfg("robot", joint_ids=slice(None)),
  )

  call_args = mock_entity.write_joint_state_to_sim.call_args
  joint_pos = call_args[0][0]
  assert torch.allclose(joint_pos, torch.ones_like(joint_pos) * 0.5)


# ===========================================================================
# Section 4: Step mode and apply_body_impulse
# ===========================================================================


def test_step_mode_fires_every_call(device):
  """Step-mode events fire unconditionally on every apply() call."""
  call_count = [0]

  def counter(env, env_ids):
    call_count[0] += 1

  env = Mock()
  env.num_envs = 2
  env.device = device
  env.scene = {}
  env.sim = Mock()

  cfg = {
    "step_counter": EventTermCfg(
      mode="step",
      func=counter,
      params={},
    ),
  }
  manager = EventManager(cfg, env)

  for _ in range(5):
    manager.apply(mode="step", dt=0.02)

  assert call_count[0] == 5


def _make_impulse_env(
  device, num_envs=2, num_bodies=1, body_ids=None, cooldown_s=(0.0, 0.0)
):
  """Create a mock env for apply_body_impulse tests."""
  if body_ids is None:
    body_ids = [0]
  env = Mock()
  env.num_envs = num_envs
  env.device = device
  env.step_dt = 0.02

  mock_entity = Mock()
  mock_entity.num_bodies = num_bodies
  mock_entity.data = Mock()
  mock_entity.data.body_com_quat_w = torch.zeros(
    (num_envs, num_bodies, 4), device=device
  )
  mock_entity.data.body_com_quat_w[..., 0] = 1.0
  env.scene = {"robot": mock_entity}

  asset_cfg = SceneEntityCfg("robot", body_ids=body_ids)
  term_cfg = Mock()
  term_cfg.params = {"asset_cfg": asset_cfg, "cooldown_s": cooldown_s}
  impulse = events.apply_body_impulse(cfg=term_cfg, env=env)
  return env, mock_entity, asset_cfg, impulse


def test_apply_body_impulse_basic(device):
  """Impulse is applied and cleared after duration expires."""
  env, mock_entity, asset_cfg, impulse = _make_impulse_env(
    device, num_envs=2, num_bodies=3, body_ids=[1], cooldown_s=(10.0, 10.0)
  )

  # Skip the initial cooldown so the first call triggers immediately;
  # the trigger/sustain/expire cycle is what's under test here.
  impulse._interval_time_left[:] = 0.0

  impulse(
    env,
    None,
    force_range=(10.0, 10.0),
    torque_range=(0.0, 0.0),
    duration_s=(0.05, 0.05),
    cooldown_s=(10.0, 10.0),
    asset_cfg=asset_cfg,
  )
  assert impulse._active.any()
  mock_entity.write_external_wrench_to_sim.assert_called()

  def step():
    impulse(
      env,
      None,
      force_range=(10.0, 10.0),
      torque_range=(0.0, 0.0),
      duration_s=(0.05, 0.05),
      cooldown_s=(10.0, 10.0),
      asset_cfg=asset_cfg,
    )

  # Step a few times (duration is 0.05s, step_dt is 0.02s).
  mock_entity.write_external_wrench_to_sim.reset_mock()
  step()  # t=0.02, remaining=0.03
  assert impulse._active.all()

  step()  # t=0.04, remaining=0.01
  assert impulse._active.all()

  step()  # t=0.06, remaining=-0.01 -> cleared
  assert not impulse._active.any()

  # Verify write_external_wrench_to_sim was called with zeros to clear.
  clear_call = None
  for call in mock_entity.write_external_wrench_to_sim.call_args_list:
    forces = call[0][0]
    if torch.all(forces == 0):
      clear_call = call
  assert clear_call is not None


def test_apply_body_impulse_with_offset(device):
  """body_point_offset adds cross-product contribution to torque."""
  env, mock_entity, asset_cfg, impulse = _make_impulse_env(
    device, num_envs=1, num_bodies=1, body_ids=[0]
  )

  impulse(
    env,
    None,
    force_range=(1.0, 1.0),
    torque_range=(0.0, 0.0),
    duration_s=(1.0, 1.0),
    cooldown_s=(0.0, 0.0),
    asset_cfg=asset_cfg,
    body_point_offset=(0.0, 0.0, 0.5),
  )

  call_args = mock_entity.write_external_wrench_to_sim.call_args
  forces = call_args[0][0]
  torques = call_args[0][1]

  # force is (1,1,1) applied at offset (0,0,0.5).
  # cross((0,0,0.5), (1,1,1)) = (-0.5, 0.5, 0)
  offset = torch.tensor([0.0, 0.0, 0.5], device=device)
  expected_extra_torque = torch.cross(offset.unsqueeze(0), forces.squeeze(0), dim=-1)
  torch.testing.assert_close(
    torques.squeeze(0), expected_extra_torque, atol=1e-5, rtol=1e-5
  )


def test_apply_body_impulse_reset_clears(device):
  """reset() zeros forces and resets internal state."""
  env, mock_entity, asset_cfg, impulse = _make_impulse_env(
    device, num_envs=2, num_bodies=1, body_ids=[0]
  )

  # Trigger impulse.
  impulse(
    env,
    None,
    force_range=(10.0, 10.0),
    torque_range=(0.0, 0.0),
    duration_s=(1.0, 1.0),
    cooldown_s=(0.0, 0.0),
    asset_cfg=asset_cfg,
  )
  assert impulse._active.all()

  # Reset env 0.
  mock_entity.write_external_wrench_to_sim.reset_mock()
  impulse.reset(env_ids=torch.tensor([0], device=device))

  assert not impulse._active[0]
  assert impulse._active[1]
  assert impulse._time_remaining[0] == 0.0

  # Verify clearing write was called for env 0.
  mock_entity.write_external_wrench_to_sim.assert_called_once()
  call_args = mock_entity.write_external_wrench_to_sim.call_args
  env_ids_arg = call_args[1]["env_ids"]
  assert len(env_ids_arg) == 1
  assert env_ids_arg[0].item() == 0


def test_apply_body_impulse_initial_cooldown(device):
  """The first call after init/reset enters cooldown, not an immediate impulse.

  Regression test for #973.
  """
  env, mock_entity, asset_cfg, impulse = _make_impulse_env(
    device, num_envs=1, num_bodies=1, body_ids=[0], cooldown_s=(0.05, 0.05)
  )

  def step():
    impulse(
      env,
      None,
      force_range=(10.0, 10.0),
      torque_range=(0.0, 0.0),
      duration_s=(1.0, 1.0),
      cooldown_s=(0.05, 0.05),  # ~2.5 steps at dt=0.02
      asset_cfg=asset_cfg,
    )

  # First two steps consume the sampled cooldown; impulse must not fire yet.
  step()
  assert not impulse._active.any()
  step()
  assert not impulse._active.any()

  # Third step crosses the cooldown boundary and triggers.
  step()
  assert impulse._active.all()

  # Reset re-enters cooldown: next step should not immediately re-trigger.
  impulse.reset(env_ids=torch.tensor([0], device=device))
  assert not impulse._active.any()
  step()
  assert not impulse._active.any()


# ===========================================================================
# Section 5: Recomputation integration
# ===========================================================================


def test_body_mass_recompute_subtreemass(device):
  """body_subtreemass is correctly recomputed after body_mass change."""
  env = create_env(device, ("body_mass", "body_subtreemass"))
  robot = env.scene["robot"]
  body_ids = robot.indexing.body_ids

  original_subtreemass = env.sim.model.body_subtreemass[:, body_ids].clone()

  env.sim.model.body_mass[0, body_ids] *= 2.0
  env.sim.recompute_constants(RecomputeLevel.set_const)

  new_subtreemass = env.sim.model.body_subtreemass[:, body_ids]

  assert not torch.allclose(new_subtreemass[0], original_subtreemass[0])
  torch.testing.assert_close(new_subtreemass[1], original_subtreemass[1])


def test_body_mass_recompute_invweight(device):
  """dof_invweight0 changes after body_mass modification."""
  env = create_env(
    device,
    ("body_mass", "body_subtreemass", "dof_invweight0", "body_invweight0"),
  )
  robot = env.scene["robot"]
  body_ids = robot.indexing.body_ids
  dof_adr = robot.indexing.joint_v_adr

  original_invweight = env.sim.model.dof_invweight0[:, dof_adr].clone()

  env.sim.model.body_mass[0, body_ids] *= 2.0
  env.sim.recompute_constants(RecomputeLevel.set_const)

  new_invweight = env.sim.model.dof_invweight0[:, dof_adr]

  assert not torch.allclose(new_invweight[0], original_invweight[0])
  torch.testing.assert_close(new_invweight[1], original_invweight[1])


def test_body_com_offset_recompute(device):
  """body_ipos modification triggers correct recomputation."""
  env = create_env(
    device,
    ("body_ipos", "body_subtreemass", "dof_invweight0", "body_invweight0"),
  )
  robot = env.scene["robot"]
  body_ids = robot.indexing.body_ids
  dof_adr = robot.indexing.joint_v_adr

  original_invweight = env.sim.model.dof_invweight0[:, dof_adr].clone()

  link1_idx = body_ids[1]
  env.sim.model.body_ipos[0, link1_idx, 0] += 0.05
  env.sim.recompute_constants(RecomputeLevel.set_const)

  new_invweight = env.sim.model.dof_invweight0[:, dof_adr]

  assert not torch.allclose(new_invweight[0], original_invweight[0])
  torch.testing.assert_close(new_invweight[1], original_invweight[1])


def test_dof_armature_recompute(device):
  """dof_armature modification updates dof_invweight0."""
  env = create_env(
    device,
    (
      "dof_armature",
      "dof_invweight0",
      "body_invweight0",
      "tendon_length0",
      "tendon_invweight0",
    ),
  )
  robot = env.scene["robot"]
  dof_adr = robot.indexing.joint_v_adr

  original_invweight = env.sim.model.dof_invweight0[:, dof_adr].clone()

  env.sim.model.dof_armature[0, dof_adr] *= 10.0
  env.sim.recompute_constants(RecomputeLevel.set_const_0)

  new_invweight = env.sim.model.dof_invweight0[:, dof_adr]

  assert not torch.allclose(new_invweight[0], original_invweight[0])
  torch.testing.assert_close(new_invweight[1], original_invweight[1])


# ===========================================================================
# Section 6: Recompute vs mj_setConst / recompiled MjSpec
# ===========================================================================


def test_body_mass_recompute_matches_set_const(device):
  """Modify body_mass: GPU recompute matches CPU mj_setConst."""
  gpu_model = mujoco.MjModel.from_xml_string(ROBOT_XML)
  cpu_model = mujoco.MjModel.from_xml_string(ROBOT_XML)

  link1_id = mujoco.mj_name2id(cpu_model, mujoco.mjtObj.mjOBJ_BODY, "link1")

  # GPU path.
  expand_fields = (
    "body_mass",
    "body_subtreemass",
    "dof_invweight0",
    "body_invweight0",
  )
  sim = Simulation(num_envs=1, cfg=SimulationCfg(), model=gpu_model, device=device)
  sim.expand_model_fields(expand_fields)
  sim.model.body_mass[0, link1_id] = 5.0
  sim.recompute_constants(RecomputeLevel.set_const)

  # CPU ground truth: apply same cfg mutations as Simulation.__init__.
  SimulationCfg().mujoco.apply(cpu_model)
  cpu_data = mujoco.MjData(cpu_model)
  mujoco.mj_forward(cpu_model, cpu_data)
  cpu_model.body_mass[link1_id] = 5.0
  mujoco.mj_setConst(cpu_model, cpu_data)

  for field in ("body_subtreemass", "dof_invweight0", "body_invweight0"):
    gpu_val = getattr(sim.model, field)[0].cpu()
    ref_val = torch.tensor(getattr(cpu_model, field), dtype=torch.float32)
    torch.testing.assert_close(gpu_val, ref_val, atol=1e-4, rtol=1e-4)


def test_body_ipos_recompute_matches_set_const(device):
  """Modify body_ipos: GPU recompute matches CPU mj_setConst."""
  gpu_model = mujoco.MjModel.from_xml_string(ROBOT_XML)
  cpu_model = mujoco.MjModel.from_xml_string(ROBOT_XML)

  link1_id = mujoco.mj_name2id(cpu_model, mujoco.mjtObj.mjOBJ_BODY, "link1")

  # GPU path.
  expand_fields = (
    "body_ipos",
    "body_subtreemass",
    "dof_invweight0",
    "body_invweight0",
  )
  sim = Simulation(num_envs=1, cfg=SimulationCfg(), model=gpu_model, device=device)
  sim.expand_model_fields(expand_fields)
  sim.model.body_ipos[0, link1_id, 0] += 0.05
  sim.recompute_constants(RecomputeLevel.set_const)

  # CPU ground truth: apply same cfg mutations as Simulation.__init__.
  SimulationCfg().mujoco.apply(cpu_model)
  cpu_data = mujoco.MjData(cpu_model)
  mujoco.mj_forward(cpu_model, cpu_data)
  cpu_model.body_ipos[link1_id, 0] += 0.05
  # Clear body_sameframe so mj_kinematics recomputes xipos from body_ipos.
  # (When sameframe != 0, mj_local2Global skips the ipos→xipos transform.)
  cpu_model.body_sameframe[link1_id] = 0
  mujoco.mj_setConst(cpu_model, cpu_data)

  for field in ("dof_invweight0", "body_invweight0"):
    gpu_val = getattr(sim.model, field)[0].cpu()
    ref_val = torch.tensor(getattr(cpu_model, field), dtype=torch.float32)
    torch.testing.assert_close(gpu_val, ref_val, atol=1e-4, rtol=1e-4)


def test_dof_armature_recompute_matches_recompile(device):
  """Modify dof_armature: GPU recompute matches fresh MjSpec.compile()."""
  spec = mujoco.MjSpec.from_string(ROBOT_XML)
  original_model = spec.compile()

  joint1_id = mujoco.mj_name2id(original_model, mujoco.mjtObj.mjOBJ_JOINT, "joint1")

  # Ground truth: modify spec armature, recompile.
  spec.joint("joint1").armature = 1.0
  recompiled_model = spec.compile()

  # GPU path.
  expand_fields = (
    "dof_armature",
    "dof_invweight0",
    "body_invweight0",
    "tendon_length0",
    "tendon_invweight0",
  )
  sim = Simulation(num_envs=1, cfg=SimulationCfg(), model=original_model, device=device)
  sim.expand_model_fields(expand_fields)

  dof_adr = original_model.jnt_dofadr[joint1_id]
  sim.model.dof_armature[0, dof_adr] = 1.0
  sim.recompute_constants(RecomputeLevel.set_const_0)

  gpu_dof_invweight = sim.model.dof_invweight0[0].cpu()
  ref_dof_invweight = torch.tensor(recompiled_model.dof_invweight0, dtype=torch.float32)
  torch.testing.assert_close(gpu_dof_invweight, ref_dof_invweight, atol=1e-4, rtol=1e-4)


# ===========================================================================
# Section: interval timer reset (issue: reset iterated the class-only list)
# ===========================================================================


def _noop_interval_event(env, env_ids):
  del env, env_ids


class _ClassIntervalEvent:
  """Class-based interval event with a reset hook."""

  def __init__(self, cfg, env):
    del cfg, env

  def reset(self, env_ids=None):
    del env_ids

  def __call__(self, env, env_ids):
    del env, env_ids


def _make_mock_env(device, num_envs=4):
  env = Mock()
  env.num_envs = num_envs
  env.device = device
  env.scene = {}
  env.sim = Mock()
  return env


def test_interval_timer_resets_for_function_terms(device):
  """Function-based interval terms get their per-env timer resampled on reset."""
  env = _make_mock_env(device)
  cfg = {
    "push": EventTermCfg(
      mode="interval", func=_noop_interval_event, interval_range_s=(5.0, 10.0)
    ),
  }
  manager = EventManager(cfg, env)

  timer = manager._interval_term_time_left[0]
  # Zero the reset env's timer so the in-range assertion below can only pass
  # if reset actually resampled it (deterministic, not clone-and-compare).
  timer[1] = 0.0
  before = timer.clone()
  reset_ids = torch.tensor([1], dtype=torch.int64, device=device)
  manager.reset(env_ids=reset_ids)

  assert 5.0 <= timer[1].item() <= 10.0
  untouched = torch.tensor([0, 2, 3], device=device)
  assert torch.equal(timer[untouched], before[untouched])


def test_interval_timer_reset_with_mixed_terms(device):
  """Mixed function and class interval terms reset the right timer slots."""
  env = _make_mock_env(device)
  cfg = {
    "push": EventTermCfg(
      mode="interval", func=_noop_interval_event, interval_range_s=(5.0, 10.0)
    ),
    "impulse": EventTermCfg(
      mode="interval", func=_ClassIntervalEvent, interval_range_s=(100.0, 200.0)
    ),
  }
  manager = EventManager(cfg, env)

  push_timer = manager._interval_term_time_left[0]
  impulse_timer = manager._interval_term_time_left[1]
  # Zero both timers for the reset env so the in-range assertions can only
  # pass if reset resampled each slot (the init draw is from the same range,
  # so comparing against it would be vacuous).
  push_timer[1] = 0.0
  impulse_timer[1] = 0.0
  reset_ids = torch.tensor([1], dtype=torch.int64, device=device)
  manager.reset(env_ids=reset_ids)

  # Each timer slot is resampled from its own term's range.
  assert 5.0 <= push_timer[1].item() <= 10.0
  assert 100.0 <= impulse_timer[1].item() <= 200.0


def test_interval_timer_reset_all_envs(device):
  """reset(env_ids=None) resamples every env's timer."""
  env = _make_mock_env(device)
  cfg = {
    "push": EventTermCfg(
      mode="interval", func=_noop_interval_event, interval_range_s=(5.0, 10.0)
    ),
  }
  manager = EventManager(cfg, env)

  timer = manager._interval_term_time_left[0]
  timer[:] = 0.0
  manager.reset(env_ids=None)
  assert ((timer >= 5.0) & (timer <= 10.0)).all()
