tests/test_physics_simulate.py

"""Tests for fpgm.physics.simulate.MujocoSimulator -- the fpgm-side subprocess plumbing
into fpgm.physics.worker_mujoco.

**Environment split, mirrored in this file's structure.** ``fpgm.physics.simulate`` itself
runs in the ``fpgm`` env (this test file's own interpreter) and needs no ``mujoco`` import
at all -- its error-handling paths (bad interpreter path, particle/space shape mismatch) are
tested unconditionally. Anything that actually *runs* the worker subprocess needs the
isolated ``mujoco`` env's interpreter to exist on this machine (see
``fpgm.physics.worker_mujoco``'s module docstring for why that env is separate from
``fpgm``'s own numpy<2 constraint) -- those tests are gated by :data:`requires_mujoco_env`
and skip cleanly, not fail, when that interpreter is absent (e.g. a CI image that only has
the ``fpgm`` env).

**The mass-independence test reproduces, in this new batched worker, the same physical
fact ``scripts/_mujoco_settle_worker.py``'s own docstring proves by hand**: a rigid body's
free-fall-and-settle trajectory under gravity + contact + friction does not depend on its
mass, because every force in the scene (gravity, contact normal force, friction) scales
with mass, so mass cancels out of F=ma. If this assertion ever failed here, that would be a
real finding about this worker's contact/mass plumbing (see :func:`_apply_theta`'s
``mj_setConst`` call in ``fpgm.physics.worker_mujoco`` -- this exact check is what caught,
during development, that a raw ``body_mass``/``body_inertia`` field edit alone leaves
contacts too soft without it) -- not something to weaken to make pass.
"""

from __future__ import annotations

from pathlib import Path

import numpy as np
import pytest
import trimesh

from fpgm.physics.simulate import DEFAULT_MUJOCO_PYTHON, MujocoSimulator
from fpgm.physics.types import (
    RIGID_PARAMS,
    BodySpec,
    HeightfieldSpec,
    ParamSpace,
    PhysicsError,
    SimSpec,
)
from fpgm.utils.timing import StepTimer

requires_mujoco_env = pytest.mark.skipif(
    not Path(DEFAULT_MUJOCO_PYTHON).exists(),
    reason=f"isolated mujoco env interpreter not found: {DEFAULT_MUJOCO_PYTHON} "
    "(set FPGM_MUJOCO_PYTHON or install the env; see fpgm.physics.worker_mujoco's docstring)",
)


# --------------------------------------------------------------------------- #
# Fixtures
# --------------------------------------------------------------------------- #


def _box_hull(tmp_path: Path, name: str = "hull.stl") -> Path:
    box = trimesh.creation.box(extents=(0.05, 0.04, 0.03))
    path = tmp_path / name
    box.export(str(path))
    return path


def _flat_heightfield(
    tmp_path: Path, *, z: float = 0.0, nx: int = 5, ny: int = 5, cell_size_m: float = 0.1,
    name: str = "hf.npz",
) -> HeightfieldSpec:
    """A flat, fully-observed ``nx x ny`` vertex grid at constant world Z -- the
    heightfield analogue of ``_free_falling_spec``'s old flat support box, centred on
    the world origin (where every test below drops its body).
    """
    height_m = np.full((ny, nx), z, dtype=np.float64)
    valid = np.ones((ny, nx), dtype=bool)
    grid_path = tmp_path / name
    np.savez(grid_path, height_m=height_m, valid=valid)
    origin_xy = np.array([-(nx - 1) * cell_size_m / 2.0, -(ny - 1) * cell_size_m / 2.0])
    return HeightfieldSpec(
        grid_path=grid_path, nx=nx, ny=ny, cell_size_m=cell_size_m, origin_xy=origin_xy,
    )


def _free_falling_spec(
    tmp_path: Path, *, n_frames: int = 10, z0: float = 0.05, support: bool = True
) -> SimSpec:
    pose = np.eye(4)
    pose[:3, 3] = [0.0, 0.0, z0]
    body = BodySpec(label="obj", mesh_path=_box_hull(tmp_path), init_pose=pose, kind="free")
    heightfield = _flat_heightfield(tmp_path, z=0.0) if support else None
    return SimSpec(
        uuid="test_drop", camera_serial="test", dt=1.0 / 30.0, n_frames=n_frames,
        bodies=[body], gripper=None, heightfield=heightfield, substeps=40,
    )


def _rigid_theta(n: int, *, log_density: np.ndarray | None = None) -> np.ndarray:
    space = ParamSpace(RIGID_PARAMS)
    theta = np.zeros((n, space.dim))
    theta[:, space.index("log_density")] = (
        log_density if log_density is not None else np.log(500.0)
    )
    theta[:, space.index("log_mu_slide")] = np.log(0.5)
    theta[:, space.index("log_mu_torsion")] = np.log(0.005)
    theta[:, space.index("log_solref_damping")] = np.log(1.0)
    theta[:, space.index("log_mu_gripper")] = np.log(0.5)
    return theta


# --------------------------------------------------------------------------- #
# Error handling -- no mujoco env needed
# --------------------------------------------------------------------------- #


def test_run_rejects_particle_space_dim_mismatch(tmp_path):
    spec = _free_falling_spec(tmp_path)
    space = ParamSpace(RIGID_PARAMS)
    bad_theta = np.zeros((3, space.dim + 1))  # one extra column
    sim = MujocoSimulator()
    with pytest.raises(PhysicsError):
        sim.run(spec, bad_theta, space)


def test_run_raises_on_missing_interpreter(tmp_path):
    spec = _free_falling_spec(tmp_path)
    space = ParamSpace(RIGID_PARAMS)
    theta = _rigid_theta(1)
    sim = MujocoSimulator(mujoco_python=Path("/no/such/interpreter"))
    with pytest.raises(PhysicsError, match="mujoco"):
        sim.run(spec, theta, space)


def test_default_mujoco_python_is_the_documented_constant():
    # FPGM_MUJOCO_PYTHON override contract -- see the module docstring's task requirement.
    # A subprocess (not importlib.reload) so this can't leave fpgm.physics.simulate's
    # module-level constant mutated for any other test in this file.
    import os
    import subprocess
    import sys

    src_dir = str(Path(__file__).resolve().parents[1] / "src")

    def _default_path(env_value: str | None) -> str:
        env = os.environ.copy()
        env["PYTHONPATH"] = src_dir
        if env_value is None:
            env.pop("FPGM_MUJOCO_PYTHON", None)
        else:
            env["FPGM_MUJOCO_PYTHON"] = env_value
        code = (
            "from fpgm.physics.simulate import DEFAULT_MUJOCO_PYTHON; "
            "print(DEFAULT_MUJOCO_PYTHON)"
        )
        out = subprocess.run(
            [sys.executable, "-c", code],
            capture_output=True, text=True, env=env, check=True,
        )
        return out.stdout.strip()

    assert _default_path(None) == "/home/quang/miniconda3/envs/mujoco/bin/python"
    assert _default_path("/custom/env/python") == "/custom/env/python"


# --------------------------------------------------------------------------- #
# Worker smoke test -- tiny 2-particle, 10-frame drop
# --------------------------------------------------------------------------- #


@requires_mujoco_env
def test_worker_smoke_free_body_falls_under_gravity(tmp_path):
    spec = _free_falling_spec(tmp_path, n_frames=10, z0=0.05, support=False)
    space = ParamSpace(RIGID_PARAMS)
    theta = _rigid_theta(2, log_density=np.log([500.0, 800.0]))

    sim = MujocoSimulator()
    result = sim.run(spec, theta, space)

    assert result.poses.shape == (2, 10, 7)
    assert np.all(result.ok), "a short, unobstructed free fall must not diverge"
    assert np.all(np.isfinite(result.poses)), "ok=True particles must have finite poses"

    z0 = result.poses[:, 0, 2]
    z_last = result.poses[:, -1, 2]
    np.testing.assert_array_almost_equal(z0, 0.05, decimal=6)
    assert np.all(z_last < z0), "the body must have fallen (z decreased) under gravity"


@requires_mujoco_env
def test_worker_smoke_quaternions_stay_normalised(tmp_path):
    spec = _free_falling_spec(tmp_path, n_frames=10, z0=0.05, support=False)
    space = ParamSpace(RIGID_PARAMS)
    theta = _rigid_theta(1)
    sim = MujocoSimulator()
    result = sim.run(spec, theta, space)
    quat_norms = np.linalg.norm(result.poses[0, :, 3:7], axis=1)
    np.testing.assert_array_almost_equal(quat_norms, 1.0, decimal=4)


# --------------------------------------------------------------------------- #
# Mass independence -- the measured physical fact this stage is built on top of
# --------------------------------------------------------------------------- #


@requires_mujoco_env
def test_mass_independence_on_settle(tmp_path):
    """Two particles differing ONLY in log_density, dropped onto a support with no
    gripper, must produce near-identical trajectories -- see module docstring.
    """
    spec = _free_falling_spec(tmp_path, n_frames=30, z0=0.05, support=True)
    space = ParamSpace(RIGID_PARAMS)
    theta = _rigid_theta(2, log_density=np.log([500.0, 5000.0]))  # 10x mass ratio

    sim = MujocoSimulator()
    result = sim.run(spec, theta, space)

    assert np.all(result.ok), (
        "both particles must settle onto the support, not diverge -- if this fails, "
        "check for tunneling (see fpgm.physics.scene.DEFAULT_SUBSTEPS's docstring) "
        "before touching the mass-independence assertion below"
    )

    pos_diff = np.abs(result.poses[0, :, :3] - result.poses[1, :, :3])
    max_pos_diff_m = float(np.max(pos_diff))
    # Real numbers observed while building this worker (2 particles, 500 vs 5000
    # kg/m^3, same box, same support): final-frame position differs by ~1e-6 m.
    # 1 mm is a generous margin -- any mass-dependence bug would blow through it by
    # orders of magnitude (the pre-mj_setConst-fix version differed by 0.14 m, see
    # fpgm.physics.worker_mujoco's module docstring).
    assert max_pos_diff_m < 1.0e-3, (
        f"mass-independence VIOLATED: max position difference {max_pos_diff_m:.6f} m "
        f"between a 500 kg/m^3 and a 5000 kg/m^3 particle on an otherwise identical "
        f"drop -- this is a real physics/plumbing finding, not noise"
    )

    z_final = result.poses[:, -1, 2]
    # Both particles must have actually settled near the support's top face (z~0),
    # not just agree with each other while both diverged the same way.
    assert np.all(np.abs(z_final) < 0.05)


@requires_mujoco_env
def test_mass_independence_free_fall_no_support(tmp_path):
    """Same check with no support at all (pure free fall) -- trivially mass-independent,
    kept as a fast sanity check distinct from the contact-engaged case above.
    """
    spec = _free_falling_spec(tmp_path, n_frames=10, z0=0.5, support=False)
    space = ParamSpace(RIGID_PARAMS)
    theta = _rigid_theta(2, log_density=np.log([100.0, 100000.0]))

    sim = MujocoSimulator()
    result = sim.run(spec, theta, space)
    assert np.all(result.ok)
    pos_diff = np.abs(result.poses[0, :, :3] - result.poses[1, :, :3])
    assert float(np.max(pos_diff)) < 1.0e-6


# --------------------------------------------------------------------------- #
# Heightfield geometry -- world metric height -> MuJoCo's [0,1] hfield contract, and
# back out again via a real settle, end to end through the actual compiled model
# (fpgm.physics.worker_mujoco._heightfield_mujoco_params's own docstring states the
# row/col-to-world-axis mapping and elevation formula as EMPIRICALLY VERIFIED against
# this MuJoCo version -- these are the tests that verify it, not a mocked/pure-math
# stand-in, since the `mujoco` package cannot be imported from the `fpgm` env this test
# file itself runs in; see that module's own docstring on the env split).
# --------------------------------------------------------------------------- #


@requires_mujoco_env
def test_worker_settles_on_heightfield_at_the_expected_metric_height(tmp_path):
    """The world->grid->MuJoCo-normalisation round trip preserves metric height: a flat
    heightfield built at world Z = 0.2 must produce a settled body at Z ~= 0.2 + (box
    half-height), not some other value implied by a sign/axis/scale error in
    ``_heightfield_mujoco_params``.
    """
    surface_z = 0.2
    box_half_height = 0.03 / 2.0  # matches _box_hull's own extents=(0.05, 0.04, 0.03)
    pose = np.eye(4)
    pose[:3, 3] = [0.0, 0.0, surface_z + 0.05]  # start just above the surface, not at z0=0.05
    body = BodySpec(label="obj", mesh_path=_box_hull(tmp_path), init_pose=pose, kind="free")
    spec = SimSpec(
        uuid="test_hfield_height", camera_serial="test", dt=1.0 / 30.0, n_frames=40,
        bodies=[body], gripper=None,
        heightfield=_flat_heightfield(tmp_path, z=surface_z), substeps=40,
    )
    space = ParamSpace(RIGID_PARAMS)
    theta = _rigid_theta(1)
    sim = MujocoSimulator()
    result = sim.run(spec, theta, space)

    assert result.ok[0], "a body dropped 5 cm onto a flat heightfield must settle, not diverge"
    z_final = float(result.poses[0, -1, 2])
    assert z_final == pytest.approx(surface_z + box_half_height, abs=5.0e-3), (
        f"settled height {z_final:.4f} disagrees with the heightfield's own world Z "
        f"({surface_z}) + box half-height ({box_half_height}) -- the metric height did "
        f"NOT round-trip through the [0,1] MuJoCo normalisation"
    )


@requires_mujoco_env
def test_worker_settles_at_elevated_heightfield_cell_not_only_flat_zero(tmp_path):
    """Same round-trip check at a NON-zero, NON-flat elevation profile (a raised block in
    the middle of an otherwise-lower grid) -- catches a bug that a flat-at-zero test alone
    could hide (e.g. `elevation_z`/`pos_z` correct only when `z_min == 0`).
    """
    nx = ny = 5
    cell = 0.1
    low_z, high_z = -0.05, 0.15
    height_m = np.full((ny, nx), low_z, dtype=np.float64)
    height_m[1:4, 1:4] = high_z  # a raised block spanning the grid's centre
    valid = np.ones((ny, nx), dtype=bool)
    grid_path = tmp_path / "raised.npz"
    np.savez(grid_path, height_m=height_m, valid=valid)
    origin_xy = np.array([-(nx - 1) * cell / 2.0, -(ny - 1) * cell / 2.0])
    heightfield = HeightfieldSpec(
        grid_path=grid_path, nx=nx, ny=ny, cell_size_m=cell, origin_xy=origin_xy,
    )

    box_half_height = 0.03 / 2.0
    pose = np.eye(4)
    pose[:3, 3] = [0.0, 0.0, high_z + 0.05]  # drop over the CENTRE, onto the raised block
    body = BodySpec(label="obj", mesh_path=_box_hull(tmp_path), init_pose=pose, kind="free")
    spec = SimSpec(
        uuid="test_hfield_raised", camera_serial="test", dt=1.0 / 30.0, n_frames=40,
        bodies=[body], gripper=None, heightfield=heightfield, substeps=40,
    )
    sim = MujocoSimulator()
    result = sim.run(spec, _rigid_theta(1), ParamSpace(RIGID_PARAMS))

    assert result.ok[0]
    z_final = float(result.poses[0, -1, 2])
    assert z_final == pytest.approx(high_z + box_half_height, abs=5.0e-3)


@requires_mujoco_env
def test_worker_never_rests_on_an_unobserved_heightfield_cell(tmp_path):
    """Unknown cells never become a fabricated resting surface: a body dropped exactly
    over a "hole" cell (surrounded by real ground at Z=0) must come to rest far BELOW
    the real ground level, not at the plausible Z=0 an interpolated/smoothed hole would
    have offered.
    """
    hole_z = -2.0  # matches the drop fpgm.physics.scene._HFIELD_HOLE_DROP_M's spirit
    # 3x3 vertex grid: every OUTER ring vertex is real ground (Z=0), the single CENTRE
    # vertex (where the body lands) is the hole. Any of the 4 surrounding cells the body
    # could sit on has 3 real corners and 1 hole corner -- bilinearly interpolated at the
    # exact centre vertex position, this returns exactly the hole's own value.
    height_m = np.zeros((3, 3), dtype=np.float64)
    height_m[1, 1] = hole_z
    valid = np.ones((3, 3), dtype=bool)
    valid[1, 1] = False  # the ONLY unobserved cell -- diagnostic-only, worker ignores it
    grid_path = tmp_path / "hole.npz"
    np.savez(grid_path, height_m=height_m, valid=valid)
    cell = 0.1
    heightfield = HeightfieldSpec(
        grid_path=grid_path, nx=3, ny=3, cell_size_m=cell, origin_xy=np.array([-cell, -cell]),
    )

    pose = np.eye(4)
    pose[:3, 3] = [0.0, 0.0, 0.1]  # dropped exactly at the grid centre (the hole)
    body = BodySpec(label="obj", mesh_path=_box_hull(tmp_path), init_pose=pose, kind="free")
    spec = SimSpec(
        uuid="test_hfield_hole", camera_serial="test", dt=1.0 / 30.0, n_frames=40,
        bodies=[body], gripper=None, heightfield=heightfield, substeps=40,
    )
    sim = MujocoSimulator()
    result = sim.run(spec, _rigid_theta(1), ParamSpace(RIGID_PARAMS))

    assert result.ok[0], "the hole is shallow enough here not to hit the numerical safety floor"
    z_final = float(result.poses[0, -1, 2])
    assert z_final < -1.0, (
        f"settled at z={z_final:.4f}, which is at or near the REAL ground (z=0) the "
        f"surrounding ring reports -- the hole cell must not have offered a plausible "
        f"resting height there"
    )


# --------------------------------------------------------------------------- #
# Divergence honesty -- reaching the safety floor must be ok=False, not a fabricated pose
# --------------------------------------------------------------------------- #


@requires_mujoco_env
def test_object_missing_the_support_diverges_honestly(tmp_path):
    """A free body that starts far above a support box (well beyond the safety-floor
    drop) must come back ok=False, with NaN poses past divergence -- never a silently
    substituted fallback pose (see fpgm.physics.worker_mujoco's module docstring).
    """
    spec = _free_falling_spec(tmp_path, n_frames=200, z0=3.0, support=True)
    space = ParamSpace(RIGID_PARAMS)
    theta = _rigid_theta(1)
    sim = MujocoSimulator()
    result = sim.run(spec, theta, space)
    assert not result.ok[0]
    assert np.any(np.isnan(result.poses[0]))


# --------------------------------------------------------------------------- #
# Timing -- None is a clean no-op, a real StepTimer records both views
# --------------------------------------------------------------------------- #


def test_run_accepts_none_timer(tmp_path):
    spec = _free_falling_spec(tmp_path)
    space = ParamSpace(RIGID_PARAMS)
    theta = _rigid_theta(1)
    sim = MujocoSimulator(mujoco_python=Path("/no/such/interpreter"))
    with pytest.raises(PhysicsError):
        sim.run(spec, theta, space, timer=None)


@requires_mujoco_env
def test_run_folds_worker_timing_into_parent_timer(tmp_path):
    spec = _free_falling_spec(tmp_path, n_frames=10, z0=0.05, support=False)
    space = ParamSpace(RIGID_PARAMS)
    theta = _rigid_theta(3)
    sim = MujocoSimulator()
    timer = StepTimer("test_timing", log_each=False)
    sim.run(spec, theta, space, timer=timer)
    summary = timer.summary()
    labels = {row["label"] for row in summary["steps"]}
    # Parent's own wall-clock view (spawn + JSON/npz I/O + subprocess) ...
    assert "simulate" in labels
    # ... and the worker's self-reported breakdown, folded in via timer.mark.
    assert "worker_step" in labels
    assert "worker_build" in labels
    assert "worker_io" in labels
    step_row = next(r for r in summary["steps"] if r["label"] == "worker_step")
    assert step_row["units"] == 3


# --------------------------------------------------------------------------- #
# simulate_batched -- chunking
# --------------------------------------------------------------------------- #


@requires_mujoco_env
def test_simulate_batched_matches_single_call_row_order(tmp_path):
    spec = _free_falling_spec(tmp_path, n_frames=10, z0=0.05, support=False)
    space = ParamSpace(RIGID_PARAMS)
    theta = _rigid_theta(5, log_density=np.log([300.0, 400.0, 500.0, 600.0, 700.0]))

    sim_single = MujocoSimulator()
    single = sim_single.run(spec, theta, space)

    sim_chunked = MujocoSimulator(max_particles_per_chunk=2)
    chunked = sim_chunked.simulate_batched(spec, theta, space)

    assert chunked.poses.shape == single.poses.shape
    assert chunked.ok.shape == single.ok.shape
    np.testing.assert_array_equal(chunked.ok, single.ok)
    # Different particle densities -> different (but each internally deterministic)
    # trajectories; chunking must not reorder or cross-contaminate rows.
    np.testing.assert_allclose(chunked.poses, single.poses, atol=1e-6)


def test_simulate_batched_below_chunk_size_is_a_single_call(tmp_path, monkeypatch):
    spec = _free_falling_spec(tmp_path)
    space = ParamSpace(RIGID_PARAMS)
    theta = _rigid_theta(2)
    sim = MujocoSimulator(mujoco_python=Path("/no/such/interpreter"), max_particles_per_chunk=100)
    calls = []

    def _fake_run(*_args, **_kwargs):
        calls.append(1)
        raise PhysicsError("x")

    monkeypatch.setattr(sim, "run", _fake_run)
    with pytest.raises(PhysicsError):
        sim.simulate_batched(spec, theta, space)
    assert calls == [1]