src/fpgm/physics/types.py

"""Stage contracts for S9: physics-parameter identification.

Mirrors the discipline of :mod:`fpgm.types` and :mod:`fpgm.datagen.types` --
every step consumes and produces a dataclass defined here, and no step reaches
into another's internals. Three of these dataclasses are also **wire formats**
(:class:`SimSpec`, :class:`SimResult`, :class:`MaterialVerdict`): they cross a
conda-env boundary as JSON, because MuJoCo and Qwen3-VL cannot be installed
into the ``fpgm`` env (its numpy is pinned ``<2`` by SAM 3.1). That is the same
subprocess-and-JSON isolation ``scripts/_mujoco_settle_worker.py`` already
established, and the reason every array here has an explicit, checked shape:
a wire format that silently accepts a transposed matrix is a bug that surfaces
three stages later as "the physics looks wrong".

**The one idea the whole stage is built around.** Parameters are inferred by
importance weighting over particles drawn *once* from the prior, with
per-episode log-likelihoods **summed**:

    log w_i = sum_e logL_e(theta_i)          w = softmax(log w)

An episode that carries no information about theta gives the same
``logL_e(theta_i)`` for every particle ``i``. Adding a constant to every entry
leaves ``softmax`` exactly unchanged. So an uninformative episode contributes
*nothing*, with no threshold, no gate, and no episode filtering anywhere in the
pipeline -- it is a property of the arithmetic, not a policy. Everything in
this module exists to keep that property true and checkable
(:class:`EpisodeLogLik` stores the raw per-particle vector, never a reduced
score, precisely so the invariant can be asserted in a test).

Corollary worth keeping in mind while reading :class:`ObservedTrack`: "did not
move" is *not* the same as "carries no information". An object that stayed put
while the gripper pushed past it rules out every theta whose simulation says it
should have slid. The arithmetic above extracts that automatically; nothing
needs to special-case it.

Array shape conventions, extending :mod:`fpgm.datagen.types`:
    T = frames on the episode's continuous *video* frame axis, N = particles,
    D = parameter-space dimension. All poses are ``(4, 4)`` homogeneous
    matrices in the ``panda_link0`` world frame -- the same frame
    ``poses.npz``'s ``T_world_obj`` is already in. Rotations crossing the wire
    are ``(w, x, y, z)`` quaternions, MuJoCo's own order, so the worker never
    has to reorder anything.
"""

from __future__ import annotations

from collections.abc import Sequence
from dataclasses import dataclass, field
from pathlib import Path
from typing import Any

import numpy as np

__all__ = [
    "PhysicsError",
    "ParamSpace",
    "RIGID_PARAMS",
    "PRISMATIC_PARAMS",
    "GaussianPrior",
    "MaterialVerdict",
    "BodySpec",
    "GripperSpec",
    "HeightfieldSpec",
    "SimSpec",
    "SimResult",
    "ObservedTrack",
    "EpisodeLogLik",
    "ParamPosterior",
    "PosteriorResult",
]


class PhysicsError(RuntimeError):
    """A physics-identification step cannot proceed on this input."""


# --------------------------------------------------------------------------- #
# Parameter space
# --------------------------------------------------------------------------- #

#: Free rigid body. Every entry is unconstrained (log or signed offset) so the
#: prior can be a plain diagonal Gaussian and particles never need rejection.
#:
#: ``log_density`` is here *despite* being expected to stay at its prior on
#: gravity-driven and quasi-static episodes -- see
#: ``scripts/_mujoco_settle_worker.py``'s docstring, which proves a rigid body's
#: fall-and-settle trajectory is mass-independent. Dropping it would hide that
#: result; keeping it makes "mass was not identifiable here" a *measured*
#: contraction of ~0 rather than an assumption.
RIGID_PARAMS: tuple[str, ...] = (
    "log_density",          # ln(kg/m^3); mass = density * mesh volume
    "com_x",                # centre-of-mass offset, units of object bounding radius
    "com_y",
    "log_mu_slide",         # ln tangential friction, object vs support
    "log_mu_torsion",       # ln torsional friction
    "log_solref_damping",   # ln contact damping; MuJoCo has no literal restitution
    "log_mu_gripper",       # ln finger-vs-object friction: what actually governs slip
)

#: Extra parameters for a body on a prismatic joint (a drawer). The joint's
#: axis and origin are **not** here: ``fpgm.datagen.events.PrismaticFit`` already
#: measures them from the observed pose track, and re-inferring a quantity that
#: was directly measured would trade a good estimate for a worse one.
PRISMATIC_PARAMS: tuple[str, ...] = (
    "log_joint_friction",
    "log_joint_damping",
)


@dataclass(frozen=True)
class ParamSpace:
    """An ordered set of unconstrained parameter names.

    Exists so that a particle array's columns have meaning that survives being
    written to disk and read back by another process. Every ``(N, D)`` array in
    this module is implicitly indexed by *some* ``ParamSpace``; passing the two
    around together is what stops a column-order mismatch from becoming a silent
    physics error.
    """

    names: tuple[str, ...]

    def __post_init__(self) -> None:
        if not self.names:
            raise PhysicsError("ParamSpace needs at least one parameter name")
        if len(set(self.names)) != len(self.names):
            raise PhysicsError(f"duplicate parameter names: {self.names}")

    @property
    def dim(self) -> int:
        return len(self.names)

    def index(self, name: str) -> int:
        try:
            return self.names.index(name)
        except ValueError as exc:
            raise PhysicsError(f"{name!r} not in ParamSpace{self.names}") from exc

    def to_dict(self, theta: np.ndarray) -> dict[str, float]:
        """One particle ``(D,)`` -> ``{name: value}``."""
        theta = np.asarray(theta, dtype=np.float64)
        if theta.shape != (self.dim,):
            raise PhysicsError(f"expected theta of shape ({self.dim},), got {theta.shape}")
        return {n: float(v) for n, v in zip(self.names, theta, strict=True)}

    def from_dict(self, values: dict[str, float]) -> np.ndarray:
        missing = [n for n in self.names if n not in values]
        if missing:
            raise PhysicsError(f"missing parameters {missing}")
        return np.array([float(values[n]) for n in self.names], dtype=np.float64)

    def extended(self, extra: Sequence[str]) -> ParamSpace:
        """This space plus ``extra`` (e.g. rigid + prismatic for a drawer)."""
        return ParamSpace(tuple(self.names) + tuple(extra))


@dataclass(frozen=True)
class GaussianPrior:
    """Diagonal Gaussian over an unconstrained :class:`ParamSpace`.

    Diagonal on purpose: the correlations that matter (mass-vs-friction and
    friction-vs-normal-load degeneracies) are properties of the *posterior*, and
    :class:`PosteriorResult`'s sensitivity SVD measures them. Baking a guessed
    correlation into the prior would make that measurement uninterpretable --
    you could no longer tell a degeneracy the data revealed from one you put in
    by hand.
    """

    space: ParamSpace
    mean: np.ndarray  # (D,)
    std: np.ndarray   # (D,), strictly positive
    provenance: dict[str, Any] = field(default_factory=dict)

    def __post_init__(self) -> None:
        d = self.space.dim
        object.__setattr__(self, "mean", np.asarray(self.mean, dtype=np.float64).reshape(-1))
        object.__setattr__(self, "std", np.asarray(self.std, dtype=np.float64).reshape(-1))
        if self.mean.shape != (d,) or self.std.shape != (d,):
            raise PhysicsError(
                f"prior mean/std must be ({d},); got {self.mean.shape}/{self.std.shape}"
            )
        if not np.all(np.isfinite(self.mean)) or not np.all(np.isfinite(self.std)):
            raise PhysicsError("prior mean/std must be finite")
        if np.any(self.std <= 0):
            raise PhysicsError(f"prior std must be > 0, got {self.std}")

    def sample(self, rng: np.random.Generator, n: int) -> np.ndarray:
        """``(n, D)`` particles drawn from the prior."""
        return self.mean[None, :] + self.std[None, :] * rng.standard_normal((n, self.space.dim))

    def logpdf(self, theta: np.ndarray) -> np.ndarray:
        """``(N,)`` log-density of ``(N, D)`` particles."""
        theta = np.atleast_2d(np.asarray(theta, dtype=np.float64))
        z = (theta - self.mean[None, :]) / self.std[None, :]
        norm = np.sum(np.log(self.std)) + 0.5 * self.space.dim * np.log(2 * np.pi)
        return -0.5 * np.sum(z * z, axis=1) - norm

    def as_dict(self) -> dict[str, Any]:
        return {
            "names": list(self.space.names),
            "mean": [float(v) for v in self.mean],
            "std": [float(v) for v in self.std],
            "provenance": self.provenance,
        }


@dataclass(frozen=True)
class MaterialVerdict:
    """A VLM's material read on one object -- a *distribution*, never a point.

    Wire format (crosses into the transformers env as JSON).

    ``classes`` is ``((material_name, probability), ...)`` and must sum to ~1.
    Keeping the full distribution rather than the argmax is the entire point: a
    VLM's material call is confident-sounding and frequently wrong, so it earns
    the right to set a *prior* and nothing more. A 60/40 wood/plastic read
    produces a genuinely wider density prior than a 99/1 read, and that width
    is what later shows up (correctly) as a wider posterior when the motion did
    not disambiguate it.
    """

    label: str
    classes: tuple[tuple[str, float], ...]
    source: str                      # e.g. "vlm:Qwen3-VL-2B-Instruct" or "table:default"
    raw: dict[str, Any] = field(default_factory=dict)

    def __post_init__(self) -> None:
        if not self.classes:
            raise PhysicsError(f"{self.label}: MaterialVerdict needs at least one class")
        total = sum(p for _, p in self.classes)
        if not np.isfinite(total) or abs(total - 1.0) > 1e-3:
            raise PhysicsError(f"{self.label}: material probabilities sum to {total}, not 1")

    def as_dict(self) -> dict[str, Any]:
        return {
            "label": self.label,
            "classes": [[str(c), float(p)] for c, p in self.classes],
            "source": self.source,
            "raw": self.raw,
        }

    @classmethod
    def from_dict(cls, d: dict[str, Any]) -> MaterialVerdict:
        return cls(
            label=str(d["label"]),
            classes=tuple((str(c), float(p)) for c, p in d["classes"]),
            source=str(d["source"]),
            raw=dict(d.get("raw", {})),
        )


# --------------------------------------------------------------------------- #
# Simulation spec -- the JSON contract with the MuJoCo worker
# --------------------------------------------------------------------------- #


@dataclass
class BodySpec:
    """One simulated body, built from artifacts S6 already produced.

    Args:
        label: The object's label in ``poses.npz`` (``"brick"``, ``"obj0"``...).
        mesh_path: Convex-hull STL written by the caller. A hull, not the raw
            mesh: MuJoCo needs convex collision geometry, and
            ``_mujoco_settle_worker.py`` already established that a hull of the
            *reconstructed* mesh beats a bounding box, because tipping depends
            on real corner geometry.
        init_pose: ``(4, 4)`` world pose at ``start_frame`` (for ``"kinematic"``,
            this is ``pose_track[0]`` -- the worker still needs a single compiled
            reference pose for the body tag even though the real driving signal is
            ``pose_track``).
        init_lin_vel / init_ang_vel: ``(3,)`` m/s and rad/s, finite-differenced
            from the observed track over the first valid frames. Unused for
            ``"kinematic"`` bodies (a mocap body has no velocity state -- its
            motion is entirely dictated by ``pose_track`` at every frame).
        kind: ``"free"``, ``"prismatic"``, or ``"kinematic"``. A ``"kinematic"``
            body is not a hypothesis being identified at all: it is another
            object S6 tracked in the same episode, driven exactly along its own
            *observed* pose track (a MuJoCo mocap body -- see
            ``fpgm.physics.worker_mujoco``'s module docstring on why mocap is
            already the right mechanism here, since it is the same one
            :class:`GripperSpec` uses for the fingers). This is what lets "the
            brick rests on the drawer" fall out of both being present with their
            measured poses, with no per-episode support-relationship logic
            anywhere in this stage.
        joint_axis_world / joint_origin_world: ``(3,)`` each, required for
            ``"prismatic"``. Measured by ``events.PrismaticFit``, not inferred.
        joint_range: ``(lo, hi)`` metres along the axis.
        pose_track: ``(T, 4, 4)`` world-frame poses, one per :class:`SimSpec`
            output frame, same convention as ``poses.npz``'s ``T_world_obj``.
            Required exactly when ``kind == "kinematic"`` (and meaningless
            otherwise, so left ``None`` for ``"free"``/``"prismatic"``).
    """

    label: str
    mesh_path: Path
    init_pose: np.ndarray
    init_lin_vel: np.ndarray = field(default_factory=lambda: np.zeros(3))
    init_ang_vel: np.ndarray = field(default_factory=lambda: np.zeros(3))
    kind: str = "free"
    joint_axis_world: np.ndarray | None = None
    joint_origin_world: np.ndarray | None = None
    joint_range: tuple[float, float] = (-1.0, 1.0)
    pose_track: np.ndarray | None = None

    def __post_init__(self) -> None:
        self.init_pose = np.asarray(self.init_pose, dtype=np.float64)
        if self.init_pose.shape != (4, 4):
            raise PhysicsError(
                f"{self.label}: init_pose must be (4, 4), got {self.init_pose.shape}"
            )
        self.init_lin_vel = np.asarray(self.init_lin_vel, dtype=np.float64).reshape(3)
        self.init_ang_vel = np.asarray(self.init_ang_vel, dtype=np.float64).reshape(3)
        if self.kind not in ("free", "prismatic", "kinematic"):
            raise PhysicsError(
                f"{self.label}: kind must be 'free', 'prismatic' or 'kinematic', "
                f"got {self.kind!r}"
            )
        if self.kind == "prismatic":
            if self.joint_axis_world is None or self.joint_origin_world is None:
                raise PhysicsError(f"{self.label}: prismatic body needs joint axis and origin")
            self.joint_axis_world = np.asarray(self.joint_axis_world, dtype=np.float64).reshape(3)
            self.joint_origin_world = np.asarray(
                self.joint_origin_world, dtype=np.float64
            ).reshape(3)
        if self.kind == "kinematic":
            if self.pose_track is None:
                raise PhysicsError(f"{self.label}: kinematic body needs a pose_track")
            self.pose_track = np.asarray(self.pose_track, dtype=np.float64)
            if self.pose_track.ndim != 3 or self.pose_track.shape[1:] != (4, 4):
                raise PhysicsError(
                    f"{self.label}: pose_track must be (T, 4, 4), got {self.pose_track.shape}"
                )


@dataclass
class GripperSpec:
    """The two fingers, as mocap bodies driven along their measured FK poses.

    This is the design choice that makes grasp, slip, push, place and release
    one model instead of five special cases: nothing in the simulation "knows"
    the object is grasped. Two finger boxes follow the poses forward kinematics
    says they had, contact does the rest, and whether the object stays in the
    hand becomes a *prediction* that depends on ``log_mu_gripper`` -- i.e. a
    thing the observed track can inform.

    Args:
        finger_poses: ``(T, 2, 4, 4)`` world poses of the two finger pads.
        finger_size: ``(3,)`` half-extents of each finger pad box, metres.
    """

    finger_poses: np.ndarray
    finger_size: np.ndarray

    def __post_init__(self) -> None:
        self.finger_poses = np.asarray(self.finger_poses, dtype=np.float64)
        if self.finger_poses.ndim != 4 or self.finger_poses.shape[1:] != (2, 4, 4):
            raise PhysicsError(
                f"finger_poses must be (T, 2, 4, 4), got {self.finger_poses.shape}"
            )
        self.finger_size = np.asarray(self.finger_size, dtype=np.float64).reshape(3)


@dataclass
class HeightfieldSpec:
    """The scene's single static collision surface: a top-down raster of the
    observed background depth, one Z per (x, y) cell.

    **Replaces the old fitted-plane :class:`SupportSpec` entirely** (see
    :mod:`fpgm.physics.scene`'s module docstring for the measured bug this fixes: a
    single fitted plane put the ONLY static geometry in the scene at the wrong
    height -- the *destination* surface, not the one the object actually starts
    on -- because a fitted plane carries no notion of "everything else in the
    room"). A heightfield built from the observed depth has geometry wherever the
    static-background camera plate actually saw something, at the height it was
    actually seen at, with no per-episode "which surface matters" judgement call.

    **Why a heightfield, not a triangle mesh or per-object boxes.** A heightfield
    needs no convex decomposition: MuJoCo's collision pipeline handles an
    ``nrow x ncol`` grid of vertices directly (bilinearly interpolated between
    them), so "cast the point cloud into *some* convex primitive per surface" --
    the problem :func:`fpgm.physics.scene._kinematic_body_spec`'s convex-hull
    export exists to solve for tracked *objects* -- never comes up for the
    static background at all. The tradeoff, stated plainly rather than glossed
    over: **a heightfield stores exactly one Z per (x, y) cell, so it cannot
    represent an overhang** -- the underside of a shelf, the inside of a closed
    drawer, a table leg's far side. Every one of those would need a second Z
    value at the same (x, y) and a heightfield structurally has none. For this
    stage's actual scenes (an object dropped/pushed/placed on top of a roughly
    horizontal work surface, observed by a downward-looking exterior camera)
    that limitation is not expected to bind -- but it is a real, permanent
    property of this representation, not a bug to be fixed later.

    **Unknown cells never become a resting surface.** A single camera leaves
    real occluded holes (the far side of the drawer, whatever the robot arm
    itself was covering in every observation). This type's contract is that a
    hole is never smoothed into a plausible-looking floor: the grid this spec
    points at is built by :func:`fpgm.physics.scene._build_heightfield_spec` so
    that (a) the grid's own XY *extent* is restricted to a region a margin
    around where this episode's tracked bodies actually were (never the whole
    room a background camera happens to see), and (b) any cell *within* that
    extent with zero directly-observed support is set to a fixed height well
    BELOW the lowest real measurement, not interpolated from its neighbours --
    see that function's docstring for the exact numbers and why interpolation
    was rejected. A particle can therefore fall through an unknown region (an
    honest, reportable divergence) but can never be handed a fabricated resting
    height there.

    **Wire format: a sidecar ``.npz``, not inline JSON arrays.** Mirrors
    :attr:`BodySpec.mesh_path`'s own precedent (a convex hull is written to a
    caller-owned scratch file and only the *path* crosses into
    :meth:`SimSpec.to_json_dict`) rather than :class:`SupportSpec`'s old style of
    a handful of plain floats inline. A heightfield sized to the demo episode's
    tabletop region is already tens of thousands of cells (`height_m` and
    `valid`, each `(ny, nx)` float64/bool) -- inlining both as JSON arrays would
    multiply the JSON payload MujocoSimulator writes to disk every particle
    batch call by a large, cell-count-dependent factor for no benefit: the two
    arrays are read verbatim by :mod:`fpgm.physics.worker_mujoco`, never
    inspected or diffed as JSON the way the small scalar fields below are.

    Args:
        grid_path: Path to an ``.npz`` with ``height_m`` (``(ny, nx)`` float64,
            world-frame Z, already hole-filled -- see above) and ``valid``
            (``(ny, nx)`` bool, ``True`` where the cell was directly observed,
            kept alongside the filled heights purely as a diagnostic/audit
            trail, not read by the worker's own physics).
        nx, ny: Grid vertex counts along X, Y (``ncol``/``nrow`` in MuJoCo's own
            ``<hfield>`` terms).
        cell_size_m: Uniform vertex spacing, metres, both axes.
        origin_xy: ``(2,)`` world ``(x, y)`` of grid vertex ``(row=0, col=0)`` --
            i.e. the grid's minimum-XY corner, not its centre.
    """

    grid_path: Path
    nx: int
    ny: int
    cell_size_m: float
    origin_xy: np.ndarray  # (2,) world frame.

    def __post_init__(self) -> None:
        self.grid_path = Path(self.grid_path)
        if self.nx < 2 or self.ny < 2:
            raise PhysicsError(
                f"HeightfieldSpec needs nx, ny >= 2 (a single vertex has no extent); "
                f"got nx={self.nx}, ny={self.ny}"
            )
        if self.cell_size_m <= 0:
            raise PhysicsError(f"HeightfieldSpec.cell_size_m must be > 0, got {self.cell_size_m}")
        self.origin_xy = np.asarray(self.origin_xy, dtype=np.float64).reshape(2)


@dataclass
class SimSpec:
    """Everything the MuJoCo worker needs, with no reference to ``fpgm``.

    Wire format. :meth:`to_json_dict` / :meth:`from_json_dict` are the only
    sanctioned way across the env boundary -- hand-building the dict on either
    side is how a field silently goes missing.
    """

    uuid: str
    camera_serial: str
    dt: float
    n_frames: int
    bodies: list[BodySpec]
    gripper: GripperSpec | None
    heightfield: HeightfieldSpec | None
    gravity: tuple[float, float, float] = (0.0, 0.0, -9.81)
    substeps: int = 4

    def __post_init__(self) -> None:
        if self.dt <= 0:
            raise PhysicsError(f"dt must be > 0, got {self.dt}")
        if self.n_frames < 2:
            raise PhysicsError(f"n_frames must be >= 2, got {self.n_frames}")
        if not self.bodies:
            raise PhysicsError("SimSpec needs at least one body")
        if self.gripper is not None and self.gripper.finger_poses.shape[0] != self.n_frames:
            raise PhysicsError(
                f"gripper has {self.gripper.finger_poses.shape[0]} frames, "
                f"SimSpec says {self.n_frames}"
            )
        for b in self.bodies:
            if b.kind == "kinematic" and b.pose_track.shape[0] != self.n_frames:  # type: ignore[union-attr]
                raise PhysicsError(
                    f"{b.label}: kinematic body has {b.pose_track.shape[0]} pose_track "  # type: ignore[union-attr]
                    f"frames, SimSpec says {self.n_frames}"
                )

    def to_json_dict(self) -> dict[str, Any]:
        def body(b: BodySpec) -> dict[str, Any]:
            out: dict[str, Any] = {
                "label": b.label,
                "mesh_path": str(b.mesh_path),
                "init_pose": b.init_pose.tolist(),
                "init_lin_vel": b.init_lin_vel.tolist(),
                "init_ang_vel": b.init_ang_vel.tolist(),
                "kind": b.kind,
                "joint_range": list(b.joint_range),
            }
            if b.kind == "prismatic":
                out["joint_axis_world"] = b.joint_axis_world.tolist()  # type: ignore[union-attr]
                out["joint_origin_world"] = b.joint_origin_world.tolist()  # type: ignore[union-attr]
            if b.kind == "kinematic":
                out["pose_track"] = b.pose_track.tolist()  # type: ignore[union-attr]
            return out

        return {
            "uuid": self.uuid,
            "camera_serial": self.camera_serial,
            "dt": self.dt,
            "n_frames": self.n_frames,
            "gravity": list(self.gravity),
            "substeps": self.substeps,
            "bodies": [body(b) for b in self.bodies],
            "gripper": None if self.gripper is None else {
                "finger_poses": self.gripper.finger_poses.tolist(),
                "finger_size": self.gripper.finger_size.tolist(),
            },
            "heightfield": None if self.heightfield is None else {
                "grid_path": str(self.heightfield.grid_path),
                "nx": self.heightfield.nx,
                "ny": self.heightfield.ny,
                "cell_size_m": self.heightfield.cell_size_m,
                "origin_xy": self.heightfield.origin_xy.tolist(),
            },
        }

    @classmethod
    def from_json_dict(cls, d: dict[str, Any]) -> SimSpec:
        bodies = [
            BodySpec(
                label=b["label"],
                mesh_path=Path(b["mesh_path"]),
                init_pose=np.array(b["init_pose"]),
                init_lin_vel=np.array(b["init_lin_vel"]),
                init_ang_vel=np.array(b["init_ang_vel"]),
                kind=b["kind"],
                joint_axis_world=(
                    np.array(b["joint_axis_world"]) if "joint_axis_world" in b else None
                ),
                joint_origin_world=(
                    np.array(b["joint_origin_world"]) if "joint_origin_world" in b else None
                ),
                joint_range=tuple(b.get("joint_range", (-1.0, 1.0))),  # type: ignore[arg-type]
                pose_track=(np.array(b["pose_track"]) if "pose_track" in b else None),
            )
            for b in d["bodies"]
        ]
        g = d.get("gripper")
        hf = d.get("heightfield")
        return cls(
            uuid=d["uuid"],
            camera_serial=d["camera_serial"],
            dt=float(d["dt"]),
            n_frames=int(d["n_frames"]),
            bodies=bodies,
            gripper=None if g is None else GripperSpec(
                finger_poses=np.array(g["finger_poses"]), finger_size=np.array(g["finger_size"])
            ),
            heightfield=None if hf is None else HeightfieldSpec(
                grid_path=Path(hf["grid_path"]),
                nx=int(hf["nx"]),
                ny=int(hf["ny"]),
                cell_size_m=float(hf["cell_size_m"]),
                origin_xy=np.array(hf["origin_xy"]),
            ),
            gravity=tuple(d.get("gravity", (0.0, 0.0, -9.81))),  # type: ignore[arg-type]
            substeps=int(d.get("substeps", 4)),
        )


@dataclass
class SimResult:
    """``(N, T, 7)`` predicted poses per particle, plus the worker's own timing.

    Wire format (the worker's reply).

    Args:
        poses: ``(N, T, 7)`` as ``[x, y, z, qw, qx, qy, qz]`` per body-of-interest,
            world frame. Quaternion order is MuJoCo's, deliberately -- see the
            module docstring.
        ok: ``(N,)`` bool. ``False`` where the rollout diverged (NaN, or the
            body reached the numerical safety floor). A diverged rollout must
            *not* be silently dropped: its particle gets ``-inf`` log-likelihood,
            which is a real statement ("this theta is inconsistent with the
            observation"), whereas dropping it would quietly reshape the
            posterior.
        label: Which body's poses these are.
        sim_seconds: Wall-clock the worker spent stepping, self-reported so the
            parent can separate MuJoCo cost from subprocess overhead.
    """

    label: str
    poses: np.ndarray
    ok: np.ndarray
    sim_seconds: float = 0.0
    n_substeps: int = 0

    def __post_init__(self) -> None:
        self.poses = np.asarray(self.poses, dtype=np.float64)
        if self.poses.ndim != 3 or self.poses.shape[2] != 7:
            raise PhysicsError(f"SimResult.poses must be (N, T, 7), got {self.poses.shape}")
        self.ok = np.asarray(self.ok, dtype=bool).reshape(-1)
        if self.ok.shape[0] != self.poses.shape[0]:
            raise PhysicsError(
                f"SimResult.ok has {self.ok.shape[0]} entries for {self.poses.shape[0]} particles"
            )


# --------------------------------------------------------------------------- #
# Observation and inference results
# --------------------------------------------------------------------------- #


@dataclass
class ObservedTrack:
    """What S6/S7 measured, plus the noise floor S6 also measured.

    ``sigma_trans_m`` comes from ``ObjectPoseDebug.pose_noise_mm`` -- the spread
    of the estimated pose over frames the object provably did not move
    (:mod:`fpgm.datagen.static_span`). Using a *measured* noise floor rather
    than a tuned one is what makes the resulting posterior calibrated instead of
    arbitrary: the likelihood's width is a property of the data, not a knob.

    Args:
        valid: ``(T,)`` bool. Frames to score. Built from ``pose_source == PNP``
            and the visibility fraction S6 already gates on -- an interpolated
            or gap-filled pose is not an observation and must not be scored as
            one.
    """

    uuid: str
    camera_serial: str
    label: str
    T_world_obj: np.ndarray  # (T, 4, 4)
    valid: np.ndarray        # (T,) bool
    sigma_trans_m: float
    sigma_rot_rad: float
    fps: float = 15.0

    def __post_init__(self) -> None:
        self.T_world_obj = np.asarray(self.T_world_obj, dtype=np.float64)
        if self.T_world_obj.ndim != 3 or self.T_world_obj.shape[1:] != (4, 4):
            raise PhysicsError(
                f"{self.label}: T_world_obj must be (T, 4, 4), got {self.T_world_obj.shape}"
            )
        self.valid = np.asarray(self.valid, dtype=bool).reshape(-1)
        if self.valid.shape[0] != self.T_world_obj.shape[0]:
            raise PhysicsError(
                f"{self.label}: valid has {self.valid.shape[0]} frames for "
                f"{self.T_world_obj.shape[0]} poses"
            )
        if not (self.sigma_trans_m > 0 and self.sigma_rot_rad > 0):
            raise PhysicsError(
                f"{self.label}: measured noise floors must be positive, got "
                f"{self.sigma_trans_m} m / {self.sigma_rot_rad} rad"
            )

    @property
    def n_valid(self) -> int:
        return int(self.valid.sum())


@dataclass
class EpisodeLogLik:
    """One episode-object's contribution: the **raw per-particle** log-likelihood.

    Stored unreduced, on purpose. The stage's central invariant -- an
    uninformative episode leaves the pooled posterior untouched -- is a
    statement about this vector being constant across ``i``, and it can only be
    asserted if the vector itself survives. A pre-reduced "score" would make the
    invariant untestable, which for a design whose entire justification *is*
    that invariant would be a poor trade.

    ``spread_nats`` (max - min over finite entries) is therefore a diagnostic,
    not a filter: ~0 means this episode said nothing about theta, and the
    correct response is to accumulate it anyway and let it be a no-op.
    """

    uuid: str
    camera_serial: str
    label: str
    loglik: np.ndarray  # (N,)
    n_obs_frames: int
    n_diverged: int = 0
    timings: dict[str, Any] = field(default_factory=dict)

    def __post_init__(self) -> None:
        self.loglik = np.asarray(self.loglik, dtype=np.float64).reshape(-1)

    @property
    def spread_nats(self) -> float:
        finite = self.loglik[np.isfinite(self.loglik)]
        return float(finite.max() - finite.min()) if finite.size else 0.0

    def as_dict(self) -> dict[str, Any]:
        return {
            "uuid": self.uuid,
            "camera_serial": self.camera_serial,
            "label": self.label,
            "n_obs_frames": self.n_obs_frames,
            "n_diverged": self.n_diverged,
            "n_particles": int(self.loglik.size),
            "spread_nats": round(self.spread_nats, 6),
            "timings": self.timings,
        }


@dataclass
class ParamPosterior:
    """Prior vs posterior for one parameter, and how much was actually learned.

    ``contraction = 1 - Var_post / Var_prior`` in [0, 1]: 0 = the data said
    nothing, 1 = the data pinned it. ``learned`` is a *label on a report*, never
    an input to a decision about which data to keep -- that distinction is the
    whole reason this stage has no filter.
    """

    name: str
    prior_mean: float
    prior_std: float
    post_mean: float
    post_std: float
    contraction: float
    learned: bool

    def as_dict(self) -> dict[str, Any]:
        return {
            "name": self.name,
            "prior": {"mean": round(self.prior_mean, 6), "std": round(self.prior_std, 6)},
            "posterior": {"mean": round(self.post_mean, 6), "std": round(self.post_std, 6)},
            "contraction": round(self.contraction, 6),
            "learned": bool(self.learned),
        }


@dataclass
class PosteriorResult:
    """The pooled posterior over one object class, plus its identifiability report.

    ``sensitivity_values`` / ``sensitivity_directions`` are the SVD of the
    weighted particle deviations: which *directions* in parameter space the data
    constrained. Contact identifiability analyses predict that mass, friction
    and applied normal load enter the observable dynamics only through certain
    combinations; this reports the combinations that were actually pinned on
    this data instead of assuming which ones they are.
    """

    space: ParamSpace
    particles: np.ndarray      # (N, D)
    log_weights: np.ndarray    # (N,) unnormalised; softmax gives the weights
    prior: GaussianPrior
    params: tuple[ParamPosterior, ...]
    ess: float
    info_nats: float
    sensitivity_values: np.ndarray      # (D,)
    sensitivity_directions: np.ndarray  # (D, D), rows are directions
    episodes: tuple[EpisodeLogLik, ...] = ()
    refinement_rounds: int = 0
    timings: dict[str, Any] = field(default_factory=dict)

    @property
    def weights(self) -> np.ndarray:
        w = self.log_weights - np.max(self.log_weights)
        w = np.exp(w)
        return w / w.sum()

    def posterior_mean(self) -> np.ndarray:
        return self.weights @ self.particles

    def as_dict(self) -> dict[str, Any]:
        return {
            "space": list(self.space.names),
            "n_particles": int(self.particles.shape[0]),
            "ess": round(float(self.ess), 3),
            "info_nats": round(float(self.info_nats), 6),
            "refinement_rounds": self.refinement_rounds,
            "prior": self.prior.as_dict(),
            "params": [p.as_dict() for p in self.params],
            "sensitivity": {
                "singular_values": [float(v) for v in self.sensitivity_values],
                "directions": [[float(x) for x in row] for row in self.sensitivity_directions],
            },
            "episodes": [e.as_dict() for e in self.episodes],
            "timings": self.timings,
        }