tests/test_run_physics_identification.py

"""Tests for scripts/run_physics_identification.py -- no GPU, no MuJoCo, no VLM, no network.

The one property this test file exists to prove, end to end: **no episode is
ever filtered**. A completely static object -- one whose simulated rollout
predicts the exact same pose regardless of which particle (physics
hypothesis) drew it -- still gets scored and pooled, and because its
per-particle log-likelihood is (up to floating point) a constant vector,
folding it into the pool changes nothing observable about the resulting
posterior. See ``scripts/run_physics_identification.py``'s own module
docstring for why that is arithmetic, not policy.

Like ``tests/test_pipeline_s9_wiring.py``, this stubs only
``fpgm.physics.scene``/``fpgm.physics.simulate`` (not written yet at the time
this test was authored) via ``sys.modules`` injection, and uses the real
``fpgm.physics.{materials,priors,likelihood,inference,report,types}`` --
including the real ``VlmPriorProposer.propose``, which raises
``PhysicsError`` (caught, falling back to ``fallback_verdict``) for a crop
path that does not exist, with no subprocess ever touched.
"""

from __future__ import annotations

import importlib.util
import json
import sys
import types
from pathlib import Path

import numpy as np
import pytest

from tests.stub_conformance import assert_stub_matches

REPO_ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(REPO_ROOT / "src"))


def _load_script():
    spec = importlib.util.spec_from_file_location(
        "run_physics_identification_under_test",
        REPO_ROOT / "scripts" / "run_physics_identification.py",
    )
    module = importlib.util.module_from_spec(spec)
    spec.loader.exec_module(module)
    return module


rpi = _load_script()

_N_FRAMES = 8
_SCENE_ID = "sceneA"
_LABEL = "brick"

# uuids whose fake simulate() ignores theta entirely and predicts the exact
# observed pose for every particle -- see module docstring.
_STATIC_UUIDS = {"ep-static"}


# --------------------------------------------------------------------------- #
# Fixture builders
# --------------------------------------------------------------------------- #


def _make_episode_dir(
    tmp_path: Path, uuid: str, camera_serial: str = "22008760"
) -> rpi.EpisodeRef:
    episode_dir = tmp_path / uuid / camera_serial
    master_dir = episode_dir / "master"
    master_dir.mkdir(parents=True, exist_ok=True)
    (episode_dir / "episode_spec.json").write_text(json.dumps({"scene_id": _SCENE_ID}))
    # Real S6 shape, just enough for EpisodePipeline-adjacent code paths that
    # might read it; run_physics_identification itself never reads poses.npz
    # content directly (fpgm.physics.scene, stubbed below, owns that).
    np.savez(
        master_dir / "poses.npz",
        labels=np.array([_LABEL]),
        **{f"{_LABEL}__pose_source": np.zeros(_N_FRAMES, dtype=np.int64)},
    )
    return rpi.EpisodeRef(uuid, camera_serial, master_dir)


def _make_fake_scene_module():
    from fpgm.physics.types import BodySpec, ObservedTrack, PhysicsError, SimSpec

    module = types.ModuleType("fpgm.physics.scene")

    def load_observed_tracks(uuid, camera_serial, *, profile=None, labels=None, timer=None):
        labels = labels if labels is not None else [_LABEL]
        poses = np.tile(np.eye(4), (_N_FRAMES, 1, 1))
        valid = np.ones(_N_FRAMES, dtype=bool)
        return {
            label: ObservedTrack(
                uuid=uuid, camera_serial=camera_serial, label=label,
                T_world_obj=poses, valid=valid,
                sigma_trans_m=0.005, sigma_rot_rad=0.02,
            )
            for label in labels
        }, {}

    def build_sim_spec(uuid, camera_serial, label, *, profile=None, scratch_dir,
                       substeps=40, n_vel_frames=5, timer=None):
        body = BodySpec(
            label=label, mesh_path=Path("dummy.stl"), init_pose=np.eye(4), kind="free",
        )
        return SimSpec(
            uuid=uuid, camera_serial=camera_serial, dt=1.0 / 15.0, n_frames=_N_FRAMES,
            bodies=[body], gripper=None, heightfield=None,
        ), {"limitations": []}

    def object_crop(uuid, camera_serial, label, *, profile=None):
        raise PhysicsError(f"stub: no crop for {label!r}")  # exercises the VLM fallback

    module.load_observed_tracks = load_observed_tracks
    module.build_sim_spec = build_sim_spec
    module.object_crop = object_crop
    assert_stub_matches("fpgm.physics.scene", module)
    return module


def _make_fake_simulate_module():
    from fpgm.physics.types import SimResult

    module = types.ModuleType("fpgm.physics.simulate")

    class MujocoSimulator:
        def __init__(self, *, mujoco_python=None, max_particles_per_chunk=512, scratch_dir=None):
            pass

        def simulate_batched(self, spec, particles, space, *, timer=None):
            n = particles.shape[0]
            body = spec.bodies[0]
            poses = np.zeros((n, spec.n_frames, 7), dtype=np.float64)
            poses[:, :, 3] = 1.0  # identity quaternion (w, x, y, z)
            if spec.uuid not in _STATIC_UUIDS:
                # Predicted x-position depends on theta's first parameter -- a
                # genuinely informative rollout: different particles disagree.
                dx = particles[:, 0] * 0.01
                poses[:, :, 0] = dx[:, None]
            ok = np.ones(n, dtype=bool)
            return SimResult(label=body.label, poses=poses, ok=ok)

    module.MujocoSimulator = MujocoSimulator
    assert_stub_matches("fpgm.physics.simulate", module)
    return module


@pytest.fixture()
def stub_physics(monkeypatch):
    monkeypatch.setitem(sys.modules, "fpgm.physics.scene", _make_fake_scene_module())
    monkeypatch.setitem(sys.modules, "fpgm.physics.simulate", _make_fake_simulate_module())
    yield


# --------------------------------------------------------------------------- #
# Tests
# --------------------------------------------------------------------------- #


class TestNoEpisodeIsEverFiltered:
    def test_including_a_static_episode_leaves_the_pooled_posterior_unchanged(
        self, tmp_path: Path, stub_physics
    ) -> None:
        moving1 = _make_episode_dir(tmp_path, "ep-moving-1")
        moving2 = _make_episode_dir(tmp_path, "ep-moving-2")
        static = _make_episode_dir(tmp_path, "ep-static")

        out_a = tmp_path / "out_a"
        out_b = tmp_path / "out_b"

        summary_a = rpi.run_physics_identification(
            [moving1, moving2], n_particles=256, seed=0, out_root=out_a,
        )
        summary_b = rpi.run_physics_identification(
            [moving1, moving2, static], n_particles=256, seed=0, out_root=out_b,
        )

        # Every episode handed in was scored -- nothing silently dropped.
        assert len(summary_a["episodes"]) == 2
        assert len(summary_b["episodes"]) == 3

        key = f"{_SCENE_ID}/{_LABEL}"
        assert key in summary_a["groups"] and key in summary_b["groups"]
        group_a, group_b = summary_a["groups"][key], summary_b["groups"][key]
        assert group_a["n_episodes"] == 2
        assert group_b["n_episodes"] == 3

        # The actual claim: the pooled POSTERIOR (not the episode list) is
        # unaffected by including the static episode.
        posterior_a = json.loads(Path(group_a["physics_posterior_json"]).read_text())
        posterior_b = json.loads(Path(group_b["physics_posterior_json"]).read_text())

        assert group_a["ess"] == pytest.approx(group_b["ess"], rel=1e-9, abs=1e-9)
        assert group_a["info_nats"] == pytest.approx(group_b["info_nats"], rel=1e-9, abs=1e-9)
        for pa, pb in zip(posterior_a["params"], posterior_b["params"], strict=True):
            assert pa["name"] == pb["name"]
            assert pa["posterior"]["mean"] == pytest.approx(pb["posterior"]["mean"], abs=1e-9)
            assert pa["posterior"]["std"] == pytest.approx(pb["posterior"]["std"], abs=1e-9)
            assert pa["contraction"] == pytest.approx(pb["contraction"], abs=1e-9)

    def test_static_episode_alone_is_informative_diagnostic_but_not_a_gate(
        self, tmp_path: Path, stub_physics, capsys
    ) -> None:
        """spread_nats is printed (a report), and is near zero for the static
        episode -- but the episode still produces a written physics.json, is
        never skipped, and is never excluded from the group."""
        static = _make_episode_dir(tmp_path, "ep-static")
        summary = rpi.run_physics_identification(
            [static], n_particles=64, seed=0, out_root=tmp_path / "out_static_only",
        )
        assert len(summary["episodes"]) == 1
        rec = summary["episodes"][0]
        assert rec["spread_nats"] == pytest.approx(0.0, abs=1e-6)
        assert Path(rec["physics_json"]).exists()

        captured = capsys.readouterr()
        assert "spread_nats=" in captured.out

    def test_no_min_motion_or_static_skip_flag_exists(self) -> None:
        """The CLI must not offer a filtering knob -- see the module docstring's
        non-negotiable design constraint. Checked against argparse's own
        recognised option strings (a real rejection of an unknown flag), not by
        scanning rendered help text, which also contains this docstring's own
        prose explaining why no such flag exists.
        """
        old_argv = sys.argv
        for bogus_flag in ("--min-motion", "--min-displacement", "--skip-static"):
            try:
                sys.argv = [
                    "run_physics_identification.py", "--uuid", "u", "--camera", "ext1",
                    bogus_flag, "0.1",
                ]
                with pytest.raises(SystemExit):
                    rpi.parse_args()
            finally:
                sys.argv = old_argv


class TestEpisodeDiscoveryFromBatch:
    def test_finds_every_episode_with_s6_output(self, tmp_path: Path) -> None:
        _make_episode_dir(tmp_path, "ep-a")
        _make_episode_dir(tmp_path, "ep-b")
        # An episode/camera with no poses.npz -- must NOT be discovered.
        no_output_dir = tmp_path / "ep-c" / "22008760" / "master"
        no_output_dir.mkdir(parents=True)

        refs = rpi.discover_episodes_from_batch(tmp_path)
        found_uuids = {r.uuid for r in refs}
        assert found_uuids == {"ep-a", "ep-b"}


class TestArgParsing:
    def test_requires_uuid_and_camera_together_or_from_batch(self) -> None:
        old_argv = sys.argv
        try:
            sys.argv = ["run_physics_identification.py", "--uuid", "u"]
            with pytest.raises(SystemExit):
                rpi.parse_args()
        finally:
            sys.argv = old_argv

    def test_from_batch_mutually_exclusive_with_uuid(self) -> None:
        old_argv = sys.argv
        try:
            sys.argv = [
                "run_physics_identification.py", "--uuid", "u", "--camera", "ext1",
                "--from-batch", "outputs/datagen",
            ]
            with pytest.raises(SystemExit):
                rpi.parse_args()
        finally:
            sys.argv = old_argv

    def test_default_particles_is_512(self) -> None:
        old_argv = sys.argv
        try:
            sys.argv = ["run_physics_identification.py", "--uuid", "u", "--camera", "ext1"]
            args = rpi.parse_args()
        finally:
            sys.argv = old_argv
        assert args.particles == 512
        assert args.seed == 0