tests/test_physics_priors.py

"""Synthetic tests for fpgm.physics.materials and fpgm.physics.priors.

No GPU, no network, no model load: :mod:`fpgm.physics.materials` is pure
arithmetic over a static table, and every :mod:`fpgm.physics.priors` test that
would otherwise need the VLM worker subprocess stubs ``subprocess.run`` and
writes its own canned JSON to the ``--out`` path the caller passed, exactly
mirroring what the real ``scripts/_vlm_material_worker.py`` would produce.
"""

from __future__ import annotations

import json
import subprocess
from pathlib import Path

import numpy as np
import pytest

from fpgm.datagen.cache import StageCache
from fpgm.physics import priors as priors_module
from fpgm.physics.materials import material_prior
from fpgm.physics.priors import VlmPriorProposer
from fpgm.physics.types import (
    PRISMATIC_PARAMS,
    RIGID_PARAMS,
    GaussianPrior,
    MaterialVerdict,
    ParamSpace,
    PhysicsError,
)


def _verdict(*class_probs: tuple[str, float], label: str = "obj") -> MaterialVerdict:
    return MaterialVerdict(label=label, classes=tuple(class_probs), source="test")


# --------------------------------------------------------------------------- #
# Moment-matched mixture
# --------------------------------------------------------------------------- #


class TestMomentMatchedMixture:
    def test_pure_class_reproduces_component(self):
        space = ParamSpace(RIGID_PARAMS)
        v_wood = _verdict(("wood", 1.0))
        v_mixed_but_degenerate = _verdict(("wood", 1.0), ("plastic", 0.0))

        p_wood = material_prior(v_wood, space)
        p_degenerate = material_prior(v_mixed_but_degenerate, space)

        di = space.index("log_density")
        # A verdict naming only wood, and one naming wood+plastic with zero
        # weight on plastic, must produce IDENTICAL priors -- the zero-weight
        # component contributes nothing to either the mean or the variance.
        assert p_wood.mean[di] == pytest.approx(p_degenerate.mean[di])
        assert p_wood.std[di] == pytest.approx(p_degenerate.std[di])

    def test_50_50_mixture_variance_exceeds_either_component(self):
        space = ParamSpace(RIGID_PARAMS)
        di = space.index("log_density")

        v_wood = _verdict(("wood", 1.0))
        v_metal = _verdict(("metal", 1.0))
        v_mixed = _verdict(("wood", 0.5), ("metal", 0.5))

        std_wood = material_prior(v_wood, space).std[di]
        std_metal = material_prior(v_metal, space).std[di]
        std_mixed = material_prior(v_mixed, space).std[di]

        # The between-class term is strictly positive whenever the two
        # medians differ (wood ~550 kg/m^3 vs metal ~4600 kg/m^3 -- they
        # differ a lot), so the 50/50 mixture must be strictly wider than
        # EITHER pure component, not just their average.
        assert std_mixed > std_wood
        assert std_mixed > std_metal

    def test_between_class_term_matches_hand_computed_formula(self):
        # Direct check of the mean/var formula against a hand-rolled
        # computation, independent of the module's own implementation, using
        # two materials with known (median, gsd) from the table.
        import math

        from fpgm.physics.materials import _DENSITY_KG_M3  # noqa: SLF001 (test-only)

        space = ParamSpace(RIGID_PARAMS)
        di = space.index("log_density")
        v = _verdict(("wood", 0.5), ("metal", 0.5))
        got = material_prior(v, space)

        med_w, gsd_w = _DENSITY_KG_M3["wood"]
        med_m, gsd_m = _DENSITY_KG_M3["metal"]
        mu_w, sig_w = math.log(med_w), math.log(gsd_w)
        mu_m, sig_m = math.log(med_m), math.log(gsd_m)

        mean_expected = 0.5 * mu_w + 0.5 * mu_m
        within = 0.5 * sig_w**2 + 0.5 * sig_m**2
        between = 0.5 * (mu_w - mean_expected) ** 2 + 0.5 * (mu_m - mean_expected) ** 2
        std_expected = math.sqrt(within + between)

        assert got.mean[di] == pytest.approx(mean_expected)
        assert got.std[di] == pytest.approx(std_expected)


# --------------------------------------------------------------------------- #
# Prior width monotonicity
# --------------------------------------------------------------------------- #


class TestPriorWidthMonotonicity:
    def test_density_std_increases_as_verdict_gets_more_ambiguous(self):
        space = ParamSpace(RIGID_PARAMS)
        di = space.index("log_density")

        ratios = [(0.99, 0.01), (0.9, 0.1), (0.75, 0.25), (0.6, 0.4), (0.5, 0.5)]
        stds = []
        for p_wood, p_plastic in ratios:
            v = _verdict(("wood", p_wood), ("plastic", p_plastic))
            stds.append(material_prior(v, space).std[di])

        assert all(b > a for a, b in zip(stds, stds[1:], strict=False)), (
            f"expected strictly increasing std as the verdict gets more ambiguous, got {stds}"
        )

    def test_three_way_ambiguity_widens_further_than_two_way(self):
        space = ParamSpace(RIGID_PARAMS)
        di = space.index("log_density")
        two_way = material_prior(_verdict(("wood", 0.5), ("plastic", 0.5)), space)
        three_way = material_prior(
            _verdict(("wood", 1 / 3), ("plastic", 1 / 3), ("metal", 1 / 3)), space
        )
        assert three_way.std[di] > two_way.std[di]


# --------------------------------------------------------------------------- #
# ParamSpace coverage
# --------------------------------------------------------------------------- #


class TestParamSpaceCoverage:
    def test_rigid_only_space(self):
        space = ParamSpace(RIGID_PARAMS)
        prior = material_prior(_verdict(("wood", 1.0)), space)
        assert prior.space.names == space.names
        assert prior.mean.shape == (space.dim,)
        assert prior.std.shape == (space.dim,)
        assert np.all(np.isfinite(prior.mean))
        assert np.all(prior.std > 0)

    def test_rigid_plus_prismatic_space(self):
        space = ParamSpace(RIGID_PARAMS).extended(PRISMATIC_PARAMS)
        prior = material_prior(_verdict(("metal", 0.7), ("plastic", 0.3)), space)
        assert prior.space.names == space.names
        assert prior.space.dim == len(RIGID_PARAMS) + len(PRISMATIC_PARAMS)
        assert np.all(np.isfinite(prior.mean))
        assert np.all(prior.std > 0)
        # Joint params are material-independent by construction: identical
        # regardless of which materials the verdict named.
        other_prior = material_prior(_verdict(("wood", 1.0)), space)
        jf = space.index("log_joint_friction")
        jd = space.index("log_joint_damping")
        assert prior.mean[jf] == pytest.approx(other_prior.mean[jf])
        assert prior.std[jf] == pytest.approx(other_prior.std[jf])
        assert prior.mean[jd] == pytest.approx(other_prior.mean[jd])
        assert prior.std[jd] == pytest.approx(other_prior.std[jd])

    def test_unknown_parameter_raises(self):
        bogus_space = ParamSpace(("log_density", "some_new_param_nobody_added_here"))
        with pytest.raises(PhysicsError):
            material_prior(_verdict(("wood", 1.0)), bogus_space)

    def test_provenance_records_derivation(self):
        space = ParamSpace(RIGID_PARAMS)
        prior = material_prior(_verdict(("wood", 0.6), ("plastic", 0.4)), space)
        assert "params" in prior.provenance
        assert "log_density" in prior.provenance["params"]
        assert prior.provenance["params"]["log_density"]["kind"] == "material_lognormal_mixture"
        assert "log_joint_friction" not in prior.provenance["params"]  # not in this space


# --------------------------------------------------------------------------- #
# MaterialVerdict's own invariant (enforced in types.py; assert it surfaces)
# --------------------------------------------------------------------------- #


class TestMaterialVerdictInvariant:
    def test_probabilities_not_summing_to_one_raises(self):
        with pytest.raises(PhysicsError):
            MaterialVerdict(
                label="obj", classes=(("wood", 0.5), ("plastic", 0.2)), source="test"
            )

    def test_empty_classes_raises(self):
        with pytest.raises(PhysicsError):
            MaterialVerdict(label="obj", classes=(), source="test")


# --------------------------------------------------------------------------- #
# Worker-output parsing (subprocess stubbed out entirely)
# --------------------------------------------------------------------------- #


class TestParseWorkerOutput:
    def test_well_formed_output_parses(self):
        blob = json.dumps(
            [
                {
                    "label": "brick",
                    "classes": [["plastic", 0.9], ["wood", 0.1]],
                    "source": "vlm:Qwen3-VL-2B-Instruct",
                    "raw": {"argmax_letter": "B"},
                }
            ]
        )
        verdicts, _ = VlmPriorProposer._parse_worker_output(blob, ["brick"])
        assert set(verdicts) == {"brick"}
        assert verdicts["brick"].source == "vlm:Qwen3-VL-2B-Instruct"

    def test_invalid_json_raises_physics_error(self):
        with pytest.raises(PhysicsError):
            VlmPriorProposer._parse_worker_output("{not valid json", ["brick"])

    def test_non_list_top_level_raises(self):
        with pytest.raises(PhysicsError):
            VlmPriorProposer._parse_worker_output(json.dumps({"label": "brick"}), ["brick"])

    def test_malformed_verdict_entry_raises_not_half_built(self):
        # Second entry is missing "classes" entirely -- a half-built parse
        # (returning the first, valid entry and silently dropping the
        # second) must not happen; the whole call must fail.
        blob = json.dumps(
            [
                {
                    "label": "brick",
                    "classes": [["plastic", 1.0]],
                    "source": "vlm:test",
                    "raw": {},
                },
                {"label": "books", "source": "vlm:test", "raw": {}},
            ]
        )
        with pytest.raises(PhysicsError):
            VlmPriorProposer._parse_worker_output(blob, ["brick", "books"])

    def test_probabilities_not_summing_to_one_raises_through_parse(self):
        blob = json.dumps(
            [{"label": "brick", "classes": [["plastic", 0.3]], "source": "vlm:test", "raw": {}}]
        )
        with pytest.raises(PhysicsError):
            VlmPriorProposer._parse_worker_output(blob, ["brick"])

    def test_missing_expected_label_raises(self):
        blob = json.dumps(
            [{"label": "brick", "classes": [["plastic", 1.0]], "source": "vlm:test", "raw": {}}]
        )
        with pytest.raises(PhysicsError):
            VlmPriorProposer._parse_worker_output(blob, ["brick", "books"])


# --------------------------------------------------------------------------- #
# VlmPriorProposer.propose, subprocess stubbed
# --------------------------------------------------------------------------- #


def _fake_verdict_json(labels: list[str]) -> str:
    return json.dumps(
        [
            {
                "label": label,
                "classes": [["wood", 0.7], ["plastic", 0.3]],
                "source": "vlm:fake",
                "raw": {"stub": True},
            }
            for label in labels
        ]
    )


class TestVlmPriorProposerPropose:
    def _make_crops(self, tmp_path: Path, labels: list[str]) -> dict[str, Path]:
        crops = {}
        for label in labels:
            p = tmp_path / f"{label}.png"
            p.write_bytes(b"\x89PNG\r\n\x1a\nfake")
            crops[label] = p
        return crops

    def test_missing_crop_raises_before_any_subprocess(self, tmp_path: Path, monkeypatch):
        called = {"n": 0}

        def fake_run(*args, **kwargs):
            called["n"] += 1
            raise AssertionError("subprocess.run must not be called")

        monkeypatch.setattr(priors_module.subprocess, "run", fake_run)
        proposer = VlmPriorProposer()
        with pytest.raises(PhysicsError):
            proposer.propose({"brick": tmp_path / "does_not_exist.png"})
        assert called["n"] == 0

    def test_propose_without_cache_calls_worker_once(self, tmp_path: Path, monkeypatch):
        crops = self._make_crops(tmp_path, ["brick", "books"])
        calls = []

        def fake_run(cmd, capture_output, text):  # noqa: ANN001, FBT002
            calls.append(cmd)
            out_idx = cmd.index("--out") + 1
            Path(cmd[out_idx]).write_text(_fake_verdict_json(["brick", "books"]))
            return subprocess.CompletedProcess(cmd, 0, stdout="", stderr="")

        monkeypatch.setattr(priors_module.subprocess, "run", fake_run)
        proposer = VlmPriorProposer(vlm_python=Path("/fake/python"))
        verdicts = proposer.propose(crops)

        assert len(calls) == 1  # one subprocess call for both misses -- batched
        assert set(verdicts) == {"brick", "books"}
        assert verdicts["brick"].source == "vlm:fake"

    def test_propose_caches_and_skips_second_call(self, tmp_path: Path, monkeypatch):
        crops = self._make_crops(tmp_path, ["brick"])
        calls = []

        def fake_run(cmd, capture_output, text):  # noqa: ANN001, FBT002
            calls.append(cmd)
            out_idx = cmd.index("--out") + 1
            Path(cmd[out_idx]).write_text(_fake_verdict_json(["brick"]))
            return subprocess.CompletedProcess(cmd, 0, stdout="", stderr="")

        monkeypatch.setattr(priors_module.subprocess, "run", fake_run)
        cache = StageCache(tmp_path / "cache_root")
        proposer = VlmPriorProposer(vlm_python=Path("/fake/python"), scene_id=8756300955)

        v1 = proposer.propose(crops, cache=cache)
        v2 = proposer.propose(crops, cache=cache)

        assert len(calls) == 1  # second propose() was a pure cache hit
        assert v1["brick"].as_dict() == v2["brick"].as_dict()

    def test_worker_nonzero_exit_raises_with_stderr_tail(self, tmp_path: Path, monkeypatch):
        crops = self._make_crops(tmp_path, ["brick"])

        def fake_run(cmd, capture_output, text):  # noqa: ANN001, FBT002
            return subprocess.CompletedProcess(cmd, 1, stdout="", stderr="boom: CUDA OOM")

        monkeypatch.setattr(priors_module.subprocess, "run", fake_run)
        proposer = VlmPriorProposer(vlm_python=Path("/fake/python"))
        with pytest.raises(PhysicsError, match="boom: CUDA OOM"):
            proposer.propose(crops)

    def test_env_var_override_used_when_no_explicit_python(self, tmp_path: Path, monkeypatch):
        crops = self._make_crops(tmp_path, ["brick"])
        seen_pythons = []

        def fake_run(cmd, capture_output, text):  # noqa: ANN001, FBT002
            seen_pythons.append(cmd[0])
            out_idx = cmd.index("--out") + 1
            Path(cmd[out_idx]).write_text(_fake_verdict_json(["brick"]))
            return subprocess.CompletedProcess(cmd, 0, stdout="", stderr="")

        monkeypatch.setattr(priors_module.subprocess, "run", fake_run)
        monkeypatch.setenv("FPGM_VLM_PYTHON", "/env/override/python")
        proposer = VlmPriorProposer()  # no explicit vlm_python
        proposer.propose(crops)
        assert seen_pythons == ["/env/override/python"]


# --------------------------------------------------------------------------- #
# Fallback verdict
# --------------------------------------------------------------------------- #


class TestFallbackVerdict:
    def test_fallback_is_unknown_at_probability_one(self):
        v = VlmPriorProposer.fallback_verdict("mystery_object")
        assert v.classes == (("unknown", 1.0),)
        assert v.label == "mystery_object"

    def test_fallback_source_distinguishable_from_real_vlm_source(self):
        fallback = VlmPriorProposer.fallback_verdict("obj")
        real = MaterialVerdict(
            label="obj", classes=(("unknown", 1.0),), source="vlm:Qwen3-VL-2B-Instruct"
        )
        # Both verdicts assign 100% to "unknown" -- indistinguishable by
        # class distribution alone, which is exactly why `source` must
        # differ and must be checked, not the distribution.
        assert fallback.classes == real.classes
        assert fallback.source != real.source
        assert fallback.source == "table:default"

    def test_fallback_feeds_material_prior_without_error(self):
        space = ParamSpace(RIGID_PARAMS)
        v = VlmPriorProposer.fallback_verdict("obj")
        prior = material_prior(v, space)
        assert isinstance(prior, GaussianPrior)
        assert np.all(np.isfinite(prior.mean))
        assert np.all(prior.std > 0)