#!/usr/bin/env python3

"""Zero-shot vs cheat-sheet: the pretraining-tax report.

Takes two agent-workflow result files — one per instruction level, never

aggregated (suite contract) — and prints the per-agent, per-subject deltas

the AGENTS.md story rests on: success rate, turns, provider tokens, help

calls, unknown-command failures, and recovery costs.

  python3 scripts/instruction_level_report.py \\

      --zero-shot results/agent-workflow/<ts>.jsonl \\

      --cheat-sheet results/agent-workflow/<ts>.jsonl

Provider-reported tokens only; rows with token_source != provider_reported

are counted in coverage but excluded from token means (null means unmeasured,

never zero).

"""

from __future__ import annotations

import argparse

import json

import statistics

from pathlib import Path

from typing import Any

def row_level(row: dict[str, Any]) -> str | None:

    level = row.get("instruction_level")

    if level is None and isinstance(row.get("agent"), dict):

        level = row["agent"].get("instruction_level")

    return str(level) if level is not None else None

def load(path: Path, expected_level: str) -> list[dict[str, Any]]:

    rows = [json.loads(line) for line in path.read_text().splitlines() if line.strip()]

    levels = {level for level in (row_level(row) for row in rows) if level}

    if levels and levels != {expected_level}:

        raise SystemExit(

            f"{path} carries instruction levels {sorted(levels)}, expected only {expected_level!r} — "

            "never mix levels in one input (suite contract)"

        )

    return rows

def agent_name(row: dict[str, Any]) -> str:

    agent = row.get("agent")

    return str(agent.get("name")) if isinstance(agent, dict) else str(agent)

def metric(rows: list[dict[str, Any]], *path: str) -> float | None:

    values = []

    for row in rows:

        node: Any = row

        for key in path:

            node = node.get(key) if isinstance(node, dict) else None

            if node is None:

                break

        if node is not None:

            values.append(float(node))

    return round(statistics.mean(values), 1) if values else None

def fmt(value: Any) -> str:

    return "unmeasured" if value is None else str(value)

def delta(zero: float | None, cheat: float | None) -> str:

    if zero is None or cheat is None or zero == 0:

        return ""

    return f"{(cheat - zero) / zero * 100:+.0f}%"

def main() -> int:

    parser = argparse.ArgumentParser(description=__doc__)

    parser.add_argument("--zero-shot", type=Path, required=True)

    parser.add_argument("--cheat-sheet", type=Path, required=True)

    args = parser.parse_args()

    zero_rows = load(args.zero_shot, "zero-shot")

    cheat_rows = load(args.cheat_sheet, "cheat-sheet")

    def grouped(rows: list[dict[str, Any]]) -> dict[tuple[str, str], list[dict[str, Any]]]:

        groups: dict[tuple[str, str], list[dict[str, Any]]] = {}

        for row in rows:

            groups.setdefault((agent_name(row), str(row.get("subject"))), []).append(row)

        return groups

    zero, cheat = grouped(zero_rows), grouped(cheat_rows)

    metrics = [

        ("turns", ("turn_metrics", "assistant_turns_total")),

        ("output tokens", ("token_metrics", "output_tokens_reported")),

        ("input tokens", ("token_metrics", "input_tokens_reported")),

        ("help calls", ("tool_call_metrics", "help_calls_total")),

        ("unknown-cmd failures", ("tool_call_metrics", "unknown_command_failures_total")),

        ("failed tool calls", ("tool_call_metrics", "failed_tool_calls_total")),

        ("turns to recovery", ("turn_metrics", "turns_to_recovery")),

        ("tokens to recovery", ("turn_metrics", "tokens_to_recovery")),

    ]

    print("# Instruction-level report (zero-shot vs cheat-sheet, never aggregated)\n")

    print("Cheat-sheet delta is relative to zero-shot; negative = the cheat sheet")

    print("removed cost. The oak-vs-git gap at zero-shot is the pretraining tax;")

    print("how much of it the cheat sheet closes is the AGENTS.md claim.\n")

    for key in sorted(set(zero) | set(cheat)):

        zero_group, cheat_group = zero.get(key, []), cheat.get(key, [])

        passes_z = sum(1 for row in zero_group if row.get("outcome") == "pass")

        passes_c = sum(1 for row in cheat_group if row.get("outcome") == "pass")

        print(f"## {key[0]} × {key[1]}")

        print(f"  success: zero-shot {passes_z}/{len(zero_group)}  cheat-sheet {passes_c}/{len(cheat_group)}")

        for label, path in metrics:

            value_z, value_c = metric(zero_group, *path), metric(cheat_group, *path)

            if value_z is None and value_c is None:

                continue

            print(f"  {label:22s} zero-shot {fmt(value_z):>10s}  cheat-sheet {fmt(value_c):>10s}  {delta(value_z, value_c)}")

        print()

    return 0

if __name__ == "__main__":

    raise SystemExit(main())