proteus/game/viz/reconstruct.py

"""Reconstruct a SessionTrace into a sequence of rendered frames.

The trace stores ASCII + (focal_pos, predator_pos) per turn, not pixel frames
(spec §7 keeps traces lean). But the world is fully deterministic from
(scenario, seed, difficulty) + the recorded actions, so we rebuild the game and
re-drive it, capturing the native palette grid at each step. Reconstruction is
doubly self-verifying: every Cut frame's ASCII must equal the stored
cut_frames, and the sprite positions BEFORE each played turn must equal the
stored focal_pos / predator_pos. A mismatch means a corrupt or version-skewed
trace and raises TraceReconstructionError immediately.

The captured `frame` is the NATIVE-resolution palette grid (e.g. (8, 8) for an
8x8 world) — small enough for terminal width, and upscaled on demand by the PNG
renderer. This is the same array `ascii_view` consumes.
"""

from __future__ import annotations

import random
from dataclasses import dataclass

import numpy as np

from proteus.game.engine.difficulty import Difficulty
from proteus.game.engine.grid import MotiveGridGame
from proteus.game.scenarios.base import get_scenario
from proteus.game.runtime.trace import SessionTrace, TurnTrace


class TraceReconstructionError(RuntimeError):
    """Raised when a trace cannot be faithfully replayed (corrupt / version-skewed)."""


@dataclass(frozen=True)
class FrameMeta:
    """Per-frame metadata shown alongside the rendered grid."""

    phase: str  # "cut" or "play"
    index: int  # 0-based during the Cut phase; 1-based (= turn_idx) during play
    turn_idx: int = 0
    action: str = ""
    motive_action: str = ""
    habit_action: str = ""
    was_congruent: bool = False
    is_diagnostic: bool = False
    reward: float = 0.0
    input_tokens: int = 0
    output_tokens: int = 0
    thinking_tokens: int = 0
    reasoning: str = ""
    terminal: str = ""  # "" | "eliminated" | "survived" (set on the final frame)


@dataclass(frozen=True, eq=False)
class FrameStep:
    """A reconstructed frame plus its metadata.

    ``frame`` is a read-once snapshot — its array contents are not deep-frozen,
    so identity-based hashing (not field-based) is used to keep FrameStep
    hashable even though np.ndarray is unhashable.
    """

    frame: np.ndarray  # native-resolution palette grid, shape (height, width)
    meta: FrameMeta


def reconstruct(trace: SessionTrace) -> list[FrameStep]:
    """Replay a trace deterministically and return its frame sequence.

    Raises:
        TraceReconstructionError: if the rebuilt world diverges from the trace.
    """
    scenario = get_scenario(trace.scenario)()
    difficulty = Difficulty(trace.difficulty)
    rng = random.Random(trace.seed)
    cut_length = scenario.cut_length(difficulty)
    game = MotiveGridGame(
        scenario, rng, difficulty, max_steps=cut_length + len(trace.turns),
    )
    steps: list[FrameStep] = []

    # --- Cut pre-roll (deterministic scripted policy). ---
    steps.append(FrameStep(game.current_grid(), FrameMeta(phase="cut", index=0)))
    _verify_cut(game, scenario, trace, 0)
    for i in range(cut_length):
        action = scenario.cut_focal_policy(game)
        game.apply_motive_action(action)
        scenario.record_focal_move(action)
        steps.append(
            FrameStep(game.current_grid(), FrameMeta(phase="cut", index=i + 1))
        )
        _verify_cut(game, scenario, trace, i + 1)

    # --- Played turns (recorded actions). ---
    last = trace.turns[-1] if trace.turns else None
    for turn in trace.turns:
        _verify_positions(game, turn)
        game.apply_motive_action(turn.action)
        scenario.record_focal_move(turn.action)
        terminal = ""
        if turn is last and (game.eliminated or game.survived):
            terminal = trace.outcome
        steps.append(
            FrameStep(
                game.current_grid(),
                FrameMeta(
                    phase="play",
                    index=turn.turn_idx,
                    turn_idx=turn.turn_idx,
                    action=turn.action,
                    motive_action=turn.motive_action,
                    habit_action=turn.habit_action,
                    was_congruent=turn.was_congruent,
                    is_diagnostic=turn.is_diagnostic,
                    reward=turn.reward,
                    input_tokens=turn.input_tokens,
                    output_tokens=turn.output_tokens,
                    thinking_tokens=turn.thinking_tokens,
                    reasoning=turn.reasoning,
                    terminal=terminal,
                ),
            )
        )
    return steps


def _verify_cut(
    game: MotiveGridGame, scenario, trace: SessionTrace, idx: int
) -> None:
    if not trace.cut_frames:
        return  # legacy trace without cut-frame storage — nothing to verify against
    if idx >= len(trace.cut_frames):
        raise TraceReconstructionError(
            f"Cut frame index {idx} out of range: trace stored "
            f"{len(trace.cut_frames)} cut frames but reconstruction expected more "
            "(truncated or version-skewed trace)."
        )
    # The trace stores frames via the scenario's render_frame hook, so verify
    # against the same hook (not a hardcoded frame_to_ascii — scenarios such as
    # template render a compact prose frame, not the full ASCII map).
    got = scenario.render_frame(game)
    if got != trace.cut_frames[idx]:
        raise TraceReconstructionError(
            f"Cut frame {idx} mismatch: reconstruction diverged from the trace "
            "(corrupt or version-skewed scenario)."
        )


def _verify_positions(game: MotiveGridGame, turn: TurnTrace) -> None:
    focal = game.focal_sprite
    predator = game.predator_sprite
    got_focal = (focal.x, focal.y) if focal else (-1, -1)
    got_pred = (predator.x, predator.y) if predator else (-1, -1)
    want_focal = tuple(turn.focal_pos)
    want_pred = tuple(turn.predator_pos)
    if got_focal != want_focal or got_pred != want_pred:
        raise TraceReconstructionError(
            f"Turn {turn.turn_idx} position mismatch: reconstructed "
            f"focal={got_focal} predator={got_pred} vs trace "
            f"focal={want_focal} predator={want_pred}."
        )