src/fpgm/physics/report.py

"""Persist a :class:`~fpgm.physics.types.PosteriorResult` and render it for a log line.

Two JSON files, both produced by the same code path
(:func:`_write_report`) because they are the same shape of thing -- a
:meth:`PosteriorResult.as_dict` dump -- at two different points in the
pipeline:

* ``physics.json`` (:func:`write_object_report`) -- one object, in one
  episode's directory, from a single :func:`~fpgm.physics.inference.accumulate`
  call. Mirrors ``poses.npz``/``events.json`` living next to the episode they
  describe.
* ``physics_posterior.json`` (:func:`write_pooled_report`) -- the
  :func:`~fpgm.physics.inference.pool`-ed result across every episode of one
  object *class*, written once per class rather than once per episode.

Both go through :func:`fpgm.utils.io.atomic_write` for the reason that module's
own docstring gives: a later stage's "is this cached?" check has no way to tell
a truncated write from a real one apart from the file simply existing, so a
process killed mid-write must never leave a partial file at the final path.

**Interface assumption, stated because it cannot be verified from this
module alone:** ``out_path`` is documented as the *directory* the episode/class
output lives in, not the file path itself -- the fixed filenames above are
appended here. This mirrors how ``poses.npz``/``events.json``/``meta.json`` are
found by name inside a known episode directory elsewhere in the datagen stages,
rather than every writer inventing its own filename. If S9's orchestrator (built
by another agent, outside this module's scope) instead wants to hand this
module a full file path, that is a one-line change at the two call sites below
-- flagged here rather than silently guessed at.
"""

from __future__ import annotations

import json
from pathlib import Path
from typing import Any

from fpgm.physics.types import PhysicsError, PosteriorResult
from fpgm.utils.io import atomic_write

__all__ = ["write_object_report", "write_pooled_report", "format_table"]

_OBJECT_FILENAME = "physics.json"
_POOLED_FILENAME = "physics_posterior.json"


def _write_report(
    result: PosteriorResult, out_dir: Path, filename: str, extra: dict[str, Any] | None
) -> Path:
    payload = result.as_dict()
    if extra:
        collision = sorted(set(extra) & set(payload))
        if collision:
            raise PhysicsError(
                f"extra keys collide with PosteriorResult.as_dict()'s own fields: {collision}"
            )
        payload = {**payload, **extra}
    data = json.dumps(payload, indent=2).encode("utf-8")
    path = Path(out_dir) / filename
    atomic_write(path, data)
    return path


def write_object_report(
    result: PosteriorResult, out_path: Path, *, extra: dict[str, Any] | None = None
) -> Path:
    """Write ``result`` as ``<out_path>/physics.json``. Returns the file path written.

    Args:
        result: A single episode-object's :func:`~fpgm.physics.inference.accumulate`
            output.
        out_path: Directory to write into (created if missing, via
            :func:`~fpgm.utils.io.atomic_write`'s own ``ensure_dir``).
        extra: Extra fields to merge into the JSON payload alongside
            :meth:`PosteriorResult.as_dict`'s own (e.g. episode uuid, camera
            serial, object label -- whatever identifies *this* report to a
            caller that only has the resulting file, not the Python object).
            Raises :class:`~fpgm.physics.types.PhysicsError` on any key
            collision rather than silently letting one side win.
    """
    return _write_report(result, Path(out_path), _OBJECT_FILENAME, extra)


def write_pooled_report(
    result: PosteriorResult, out_path: Path, *, extra: dict[str, Any] | None = None
) -> Path:
    """Write ``result`` as ``<out_path>/physics_posterior.json``. Returns the file path written.

    Same payload shape as :func:`write_object_report`; different filename and
    call site (once per object class, after :func:`~fpgm.physics.inference.pool`,
    rather than once per episode).
    """
    return _write_report(result, Path(out_path), _POOLED_FILENAME, extra)


def format_table(result: PosteriorResult) -> str:
    """Human-readable per-parameter table for a log line: prior/posterior/contraction/learned.

    Rows follow ``result.space.names`` order, not a "most-informative first"
    sort: this table exists to compare one parameter across many logged runs at
    a glance, which only works if that parameter is always on the same row.
    ``result.params`` is expected to already be built in that order (both
    :func:`~fpgm.physics.inference.accumulate` and
    :func:`~fpgm.physics.inference.pool` construct it by iterating
    ``prior.space.names``), but this function re-indexes by name explicitly
    rather than trusting positional order -- printing one parameter's numbers
    under another's name would be a silent, easy-to-miss bug in exactly the
    kind of log line nobody double-checks against the source of truth.
    """
    by_name = {p.name: p for p in result.params}
    missing = [n for n in result.space.names if n not in by_name]
    if missing:
        raise PhysicsError(f"format_table: PosteriorResult has no params for {missing}")

    name_w = max(len(n) for n in result.space.names)
    header = (
        f"{'param':<{name_w}}  {'prior (mean +- std)':>22}  {'posterior (mean +- std)':>24}  "
        f"{'contract':>8}  learned"
    )
    lines = [header, "-" * len(header)]
    for name in result.space.names:
        p = by_name[name]
        prior_s = f"{p.prior_mean:+.4f} +- {p.prior_std:.4f}"
        post_s = f"{p.post_mean:+.4f} +- {p.post_std:.4f}"
        lines.append(
            f"{name:<{name_w}}  {prior_s:>22}  {post_s:>24}  "
            f"{p.contraction:>8.3f}  {'yes' if p.learned else 'no'}"
        )
    return "\n".join(lines)