Blame for tiny_computer_agent/trace.py

"""Traces: an episode you can save, replay and score.

An :class:`~tiny_computer_agent.loop.Episode` is the record of one run; a **trace**
is that record made portable. A trace keeps each frame's *element index* - the
labels and their geometry - but never the pixels, so it replays on a machine with
no display, no OCR engine and no model weights: hand the recorded frames back to a
brain and count how often it picks the element the trace says was right.

That count is the whole point of this module. It is what turns "the model looked
plausible on that one screenshot" into an **accuracy** you can compare between
brains, put in a table, and regression-test in CI - the harness's own answer to the
question the results table keeps asking: *which brain is actually right, and how
often?*

Two judgements are kept apart on purpose. The score is about the brain's
**decision** - what it chose when it had to commit. Whether the harness *allowed*
that decision to run is a separate axis, reported by :attr:`Verdict.blocked`, so a
policy refusal can never masquerade as a clean success (the harness owns success)
nor hide a brain that chose correctly.
"""

from __future__ import annotations

import json
from dataclasses import dataclass, field, replace
from pathlib import Path
from typing import Any, Callable, Mapping, Sequence

from tiny_computer_agent.actions import Action, parse_action
from tiny_computer_agent.brain import Brain
from tiny_computer_agent.loop import AgentLoop, Episode
from tiny_computer_agent.observe import (
    Element,
    FakeScreen,
    Observation,
    normalise_label,
)
from tiny_computer_agent.safety import Policy

TRACE_VERSION = 1
DONE = "done"
NOTHING = "nothing"

# The fields of Element, in one place, so serialisation and Element stay in step.
_ELEMENT_FIELDS = ("label", "kind", "x", "y", "width", "height")


def _element_to_dict(element: Element) -> dict[str, Any]:
    return {name: getattr(element, name) for name in _ELEMENT_FIELDS}


def _element_from_dict(data: Mapping[str, Any]) -> Element:
    return Element(**{name: data[name] for name in _ELEMENT_FIELDS if name in data})


@dataclass(frozen=True)
class TraceStep:
    """One recorded frame: what was on screen, and what the brain did about it."""

    index: int
    elements: tuple[Element, ...] = ()
    text: str = ""
    action: Action | None = None
    error: str | None = None
    outcome: str = ""

    def observation(self) -> Observation:
        """The frame as the brain saw it - pixels deliberately omitted."""
        return Observation(index=self.index, elements=self.elements, text=self.text)

    def to_dict(self) -> dict[str, Any]:
        return {
            "index": self.index,
            "elements": [_element_to_dict(e) for e in self.elements],
            "text": self.text,
            "action": self.action.to_dict() if self.action is not None else None,
            "error": self.error,
            "outcome": self.outcome,
        }

    @classmethod
    def from_dict(cls, data: Mapping[str, Any]) -> TraceStep:
        raw_action = data.get("action")
        return cls(
            index=int(data.get("index", 0)),
            elements=tuple(_element_from_dict(e) for e in data.get("elements", ())),
            text=data.get("text", ""),
            action=parse_action(raw_action) if raw_action else None,
            error=data.get("error"),
            outcome=data.get("outcome", ""),
        )


@dataclass
class Trace:
    """A saved run: the frames, the decisions, and the answer that was right.

    ``expected`` is the ground truth a scorer compares against. It lives on the
    trace rather than on a step because a task has one correct *choice*, not one
    correct action per turn - the waits and retries on the way there are not
    decisions about what to do.
    """

    goal: str
    expected: str | None = None
    brain: str = ""
    scenario: str = ""
    steps: list[TraceStep] = field(default_factory=list)
    success: bool = False
    stop_reason: str = ""

    # -- recording ---------------------------------------------------------

    @classmethod
    def record(
        cls,
        episode: Episode,
        *,
        expected: str | None = None,
        brain: str = "",
        scenario: str = "",
    ) -> Trace:
        """Copy an episode into a portable trace (no pixels, no live objects)."""
        return cls(
            goal=episode.goal,
            expected=expected,
            brain=brain,
            scenario=scenario,
            steps=[
                TraceStep(
                    index=step.index,
                    elements=tuple(step.observation.elements),
                    text=step.observation.text,
                    action=step.action,
                    error=step.error,
                    outcome=step.result.summary() if step.result is not None else "",
                )
                for step in episode.steps
            ],
            success=episode.success,
            stop_reason=episode.stop_reason,
        )

    # -- reading -----------------------------------------------------------

    @property
    def name(self) -> str:
        return self.scenario or self.goal

    @property
    def frames(self) -> int:
        return len(self.steps)

    def decision_step(self) -> TraceStep | None:
        """The step where the brain committed: first labelled click, else first ``done``.

        Waits are skipped - waiting is patience, not a choice about what to do.
        """
        for step in self.steps:
            action = step.action
            if action is None:
                continue
            if action.kind == "click" and action.label:
                return step
            if action.kind == "done":
                return step
        return None

    @property
    def choice(self) -> str | None:
        """The label the brain chose, or None when it committed to no element."""
        step = self.decision_step()
        if step is not None and step.action is not None and step.action.kind == "click":
            return step.action.label
        return None

    @property
    def decided(self) -> str:
        """What the brain committed to: a label, the bare word ``done``, or nothing.

        Kept deliberately blunt so a mis-scored run shows up in the table: a brain
        that answered "done" before ever looking has not chosen the right element.
        """
        step = self.decision_step()
        if step is None or step.action is None:
            return NOTHING
        return step.action.label if step.action.kind == "click" else DONE

    @property
    def blocked(self) -> bool:
        """True when the harness stopped the deciding action (e.g. a policy refusal)."""
        step = self.decision_step()
        return step is not None and step.error is not None

    def observations(self) -> list[Observation]:
        return [step.observation() for step in self.steps]

    def render(self) -> str:
        """A compact, replayable summary - the trace's own version of Episode.render()."""
        header = f"trace: {self.name}"
        if self.brain:
            header += f"  (brain: {self.brain})"
        lines = [
            header,
            f"goal: {self.goal}",
            f"expected: {self.expected!r}   chose: {self.decided!r}"
            + ("   (the harness refused it)" if self.blocked else ""),
            f"result: {'success' if self.success else 'failure'} - {self.stop_reason}",
        ]
        for step in self.steps:
            if step.action is not None:
                body = step.action.describe()
            else:
                body = step.error or "no action"
            lines.append(f"  {step.index} | {body}")
        return "\n".join(lines)

    # -- replay ------------------------------------------------------------

    def replay(self, brain: Brain) -> Trace:
        """Answer the recorded frames again with a (usually different) brain.

        Nothing is captured and nothing is clicked: this is the offline half of
        "replay a trace", and it is why a trace can be scored on any machine.
        """
        steps: list[TraceStep] = []
        for step in self.steps:
            try:
                action = brain.decide(self.goal, step.observation())
                steps.append(replace(step, action=action, error=None, outcome="replayed"))
            except Exception as exc:  # a brain that fails is a scored outcome, not a crash
                steps.append(
                    replace(step, action=None, error=f"{type(exc).__name__}: {exc}", outcome="")
                )
        return Trace(
            goal=self.goal,
            expected=self.expected,
            brain=getattr(brain, "name", type(brain).__name__),
            scenario=self.scenario,
            steps=steps,
            success=False,
            stop_reason="replayed",
        )

    # -- persistence -------------------------------------------------------

    def to_dict(self) -> dict[str, Any]:
        return {
            "trace_version": TRACE_VERSION,
            "goal": self.goal,
            "expected": self.expected,
            "brain": self.brain,
            "scenario": self.scenario,
            "success": self.success,
            "stop_reason": self.stop_reason,
            "steps": [step.to_dict() for step in self.steps],
        }

    @classmethod
    def from_dict(cls, data: Mapping[str, Any]) -> Trace:
        version = int(data.get("trace_version", TRACE_VERSION))
        if version > TRACE_VERSION:
            raise ValueError(
                f"trace_version {version} is newer than this build understands "
                f"({TRACE_VERSION}); upgrade tiny-computer-agent to read it"
            )
        return cls(
            goal=data.get("goal", ""),
            expected=data.get("expected"),
            brain=data.get("brain", ""),
            scenario=data.get("scenario", ""),
            steps=[TraceStep.from_dict(step) for step in data.get("steps", ())],
            success=bool(data.get("success", False)),
            stop_reason=data.get("stop_reason", ""),
        )

    def to_json(self, *, indent: int = 2) -> str:
        return json.dumps(self.to_dict(), indent=indent, ensure_ascii=False)

    @classmethod
    def from_json(cls, text: str) -> Trace:
        return cls.from_dict(json.loads(text))

    def save(self, path: str | Path) -> Path:
        target = Path(path)
        target.parent.mkdir(parents=True, exist_ok=True)
        target.write_text(self.to_json() + "\n", encoding="utf-8")
        return target

    @classmethod
    def load(cls, path: str | Path) -> Trace:
        return cls.from_json(Path(path).read_text(encoding="utf-8"))


def load_traces(path: str | Path) -> list[Trace]:
    """Every trace in a file or a directory of ``*.json`` traces, in name order."""
    target = Path(path)
    files = sorted(target.glob("*.json")) if target.is_dir() else [target]
    return [Trace.load(file) for file in files]


# -- scoring ---------------------------------------------------------------


@dataclass(frozen=True)
class Verdict:
    """One trace, judged: did the brain choose what the trace says was right?"""

    scenario: str
    expected: str | None
    chosen: str
    correct: bool
    brain: str = ""
    blocked: bool = False
    detail: str = ""

    def line(self) -> str:
        mark = "correct  " if self.correct else "incorrect"
        line = f"  {mark} {self.scenario}: expected {self.expected!r}, chose {self.chosen!r}"
        if self.blocked:
            line += " (harness refused it)"
        elif self.detail:
            line += f" ({self.detail})"
        return line


def score_trace(trace: Trace) -> Verdict:
    """Judge one recorded choice against the trace's ground truth.

    Comparison is on :func:`~tiny_computer_agent.observe.normalise_label`, so a
    choice the OCR rendered as ``"Save ."`` is still credited for ``"Save"`` - the
    same forgiveness the element index uses everywhere else. A decision the policy
    refused is still scored on the *choice* (that is what the benchmark measures),
    and flagged by :attr:`Verdict.blocked` rather than counted as a clean success.
    """
    chosen = trace.decided
    expected = trace.expected
    if expected is None:
        return Verdict(
            scenario=trace.name, expected=None, chosen=chosen, correct=False,
            brain=trace.brain, blocked=trace.blocked,
            detail="trace records no expected answer to score against",
        )
    correct = normalise_label(chosen) == normalise_label(expected)
    detail = "" if correct else (
        f"chose {chosen!r}, not {expected!r}" if chosen != NOTHING
        else "never committed to a choice"
    )
    return Verdict(
        scenario=trace.name, expected=expected, chosen=chosen, correct=correct,
        brain=trace.brain, blocked=trace.blocked, detail=detail,
    )


def score_replay(trace: Trace, brain: Brain) -> Verdict:
    """"Replay a trace, count how often the choice was right", for a single trace."""
    return score_trace(trace.replay(brain))


@dataclass
class Scoreboard:
    """Verdicts collected into an accuracy - the number this whole module exists for."""

    verdicts: list[Verdict] = field(default_factory=list)

    def __len__(self) -> int:
        return len(self.verdicts)

    def __iter__(self):
        return iter(self.verdicts)

    def add(self, verdict: Verdict) -> Scoreboard:
        self.verdicts.append(verdict)
        return self

    @property
    def total(self) -> int:
        return len(self.verdicts)

    @property
    def earned(self) -> int:
        return sum(1 for verdict in self.verdicts if verdict.correct)

    @property
    def accuracy(self) -> float:
        """Share of decisions that matched the right answer, 0.0 when empty."""
        return self.earned / self.total if self.total else 0.0

    @property
    def blocked(self) -> int:
        """Decisions the harness stopped, whatever the score - a separate axis."""
        return sum(1 for verdict in self.verdicts if verdict.blocked)

    def brains(self) -> list[str]:
        """Every brain that appears, in first-seen order, so tables stay stable."""
        seen: list[str] = []
        for verdict in self.verdicts:
            if verdict.brain not in seen:
                seen.append(verdict.brain)
        return seen

    def accuracy_for(self, brain: str) -> float:
        rows = [v for v in self.verdicts if v.brain == brain]
        return sum(1 for v in rows if v.correct) / len(rows) if rows else 0.0

    def by_brain(self) -> dict[str, float]:
        return {brain: self.accuracy_for(brain) for brain in self.brains()}

    def to_rows(self) -> list[dict[str, Any]]:
        """Rows for a dataframe / notebook table. ``correct`` is a plain bool."""
        return [
            {
                "brain": verdict.brain or "-",
                "scenario": verdict.scenario,
                "expected": verdict.expected,
                "chose": verdict.chosen,
                "correct": verdict.correct,
            }
            for verdict in self.verdicts
        ]

    def confusion(self) -> dict[tuple[str, str], int]:
        """(expected, chose) -> count, for the cases a scalar accuracy hides."""
        counts: dict[tuple[str, str], int] = {}
        for verdict in self.verdicts:
            key = (verdict.expected or "-", verdict.chosen)
            counts[key] = counts.get(key, 0) + 1
        return counts

    def render(self) -> str:
        lines = [f"scoreboard: {self.earned}/{self.total} correct ({self.accuracy:.0%})"]
        lines += [f"  {brain or '-'}: {self.accuracy_for(brain):.0%}" for brain in self.brains()]
        lines += [verdict.line() for verdict in self.verdicts]
        if self.blocked:
            lines.append(f"  ({self.blocked} decision(s) stopped by the harness)")
        return "\n".join(lines)


# -- a suite of tasks, for benchmarking -----------------------------------


@dataclass(frozen=True)
class Scenario:
    """A task with one element that is the right answer, and the frames to show."""

    goal: str
    frames: tuple[tuple[Element, ...], ...]
    expected: str
    text: str = ""
    name: str = ""

    def __post_init__(self) -> None:
        if not self.frames:
            raise ValueError("a scenario needs at least one frame")
        # freeze the frames so a scenario can be replayed and compared safely
        object.__setattr__(self, "frames", tuple(tuple(frame) for frame in self.frames))

    @property
    def label(self) -> str:
        return self.name or self.goal

    def screen(self) -> FakeScreen:
        """A fresh screen each time, so one scenario can be run against many brains."""
        return FakeScreen(list(self.frames), text=self.text)


def run_scenario(
    scenario: Scenario, brain: Brain, *, brain_name: str = "", max_steps: int = 6
) -> Trace:
    """Run one scenario and return its trace, ground truth included."""
    episode = AgentLoop(
        screen=scenario.screen(),
        brain=brain,
        policy=Policy(step_budget=max_steps),
    ).run(scenario.goal)
    return Trace.record(
        episode,
        expected=scenario.expected,
        brain=brain_name or getattr(brain, "name", type(brain).__name__),
        scenario=scenario.label,
    )


def score(scenario: Scenario, brain: Brain, *, brain_name: str = "", max_steps: int = 6) -> Verdict:
    return score_trace(run_scenario(scenario, brain, brain_name=brain_name, max_steps=max_steps))


def benchmark(
    scenarios: Sequence[Scenario],
    brains: Mapping[str, Callable[[Scenario], Brain]],
    *,
    max_steps: int = 6,
) -> Scoreboard:
    """Run every scenario through every brain, scoring the choice each one made.

    ``brains`` maps a display name to a *factory* taking the scenario, because some
    brains (a scripted :class:`~tiny_computer_agent.brain.MockBrain`, say) need to see
    the task before they can answer it, and because a fresh brain per task keeps state
    like "already clicked this" from leaking between scenarios.
    """
    board = Scoreboard()
    for scenario in scenarios:
        for name, factory in brains.items():
            board.add(score(scenario, factory(scenario), brain_name=name, max_steps=max_steps))
    return board


def benchmark_traces(traces: Sequence[Trace], brain: Brain, *, name: str = "") -> Scoreboard:
    """Replay recorded traces against one brain and count how often it was right.

    The literal roadmap step: the frames are already recorded, nothing new is
    captured, and the only variable is the brain under test.
    """
    label = name or getattr(brain, "name", type(brain).__name__)
    board = Scoreboard()
    for trace in traces:
        replayed = trace.replay(brain)
        replayed.brain = label
        board.add(score_trace(replayed))
    return board


__all__ = [
    "DONE",
    "NOTHING",
    "TRACE_VERSION",
    "Scenario",
    "Scoreboard",
    "Trace",
    "TraceStep",
    "Verdict",
    "benchmark",
    "benchmark_traces",
    "load_traces",
    "run_scenario",
    "score",
    "score_replay",
    "score_trace",
]