proteus/cli.py

"""Command-line interface for the PROTEUS arena.

Subcommands:
    run             Run one session and append its trace to a JSONL file.
    list-scenarios  List the registered scenario names.
    replay          Print a saved trace's turns, outcome, and metrics.
    compare         Aggregate traces by (model, difficulty) for baseline comparison.

All commands are network-free except ``run`` with a real provider spec; use
``--model fake:<name>`` for an offline smoke.
"""

from __future__ import annotations

import argparse
import sys

import proteus.grid  # noqa: F401  — side-effect: populate the scenario registry
from proteus.agents import HumanAgent, VanillaAgent
from proteus.grid.difficulty import Difficulty
from proteus.grid.scenario import list_scenarios
from proteus.providers import available_providers, make_provider
from proteus.runtime import SessionRunner, append_trace, read_traces


def _cmd_run(args: argparse.Namespace) -> int:
    try:
        provider = make_provider(args.model)
    except ValueError as exc:
        print(str(exc), file=sys.stderr)
        return 2
    if args.scenario not in list_scenarios():
        print(
            f"Unknown scenario {args.scenario!r}. "
            f"Available scenarios: {', '.join(list_scenarios())}.",
            file=sys.stderr,
        )
        return 2
    agent = VanillaAgent(provider)
    runner = SessionRunner(
        args.scenario,
        agent,
        difficulty=Difficulty(args.difficulty),
        seed=args.seed,
        play_turns=args.play_turns,
        use_probe=not args.no_probe,
    )
    trace = runner.run()
    written = append_trace(trace, args.out)
    accuracy = trace.metrics["motive_reading_accuracy"]
    print(
        f"{trace.scenario} seed={trace.seed} {trace.difficulty} "
        f"model={trace.model} -> {trace.outcome} "
        f"| motive_reading_accuracy={accuracy:.1f}% "
        f"| reactivity_index={trace.metrics['reactivity_index']:.1f}%"
    )
    print(f"trace appended to {written}")
    return 0


def _cmd_play(args: argparse.Namespace) -> int:
    if args.scenario not in list_scenarios():
        print(
            f"Unknown scenario {args.scenario!r}. "
            f"Available scenarios: {', '.join(list_scenarios())}.",
            file=sys.stderr,
        )
        return 2
    agent = HumanAgent()
    runner = SessionRunner(
        args.scenario,
        agent,
        difficulty=Difficulty(args.difficulty),
        seed=args.seed,
        play_turns=args.play_turns,
        use_probe=args.probe,
    )
    try:
        trace = runner.run()
    except EOFError:
        print("stdin closed before the session finished.", file=sys.stderr)
        return 2
    accuracy = trace.metrics["motive_reading_accuracy"]
    print(
        f"{trace.scenario} seed={trace.seed} {trace.difficulty} "
        f"model={trace.model} -> {trace.outcome} "
        f"| motive_reading_accuracy={accuracy:.1f}% "
        f"| reactivity_index={trace.metrics['reactivity_index']:.1f}%"
    )
    if args.out:
        written = append_trace(trace, args.out)
        print(f"trace appended to {written}")
    return 0


def _cmd_list_scenarios(_args: argparse.Namespace) -> int:
    for name in list_scenarios():
        print(name)
    return 0


def _cmd_replay(args: argparse.Namespace) -> int:
    try:
        traces = read_traces(args.trace_file)
    except FileNotFoundError:
        print(f"trace file not found: {args.trace_file}", file=sys.stderr)
        return 2
    if not traces:
        print(f"no traces in {args.trace_file}", file=sys.stderr)
        return 1

    if args.visual or args.png:
        from proteus.viz import reconstruct, render_terminal, write_pngs

        for trace in traces:
            steps = reconstruct(trace)
            if args.visual:
                render_terminal(steps, fps=args.fps)
            if args.png:
                paths = write_pngs(steps, args.png)
                print(f"wrote {len(paths)} PNG frames to {args.png}")
        return 0

    for trace in traces:
        print(
            f"=== {trace.scenario} seed={trace.seed} {trace.difficulty} "
            f"model={trace.model} -> {trace.outcome} ==="
        )
        for turn in trace.turns:
            tag = "congruent" if turn.was_congruent else "DIVERGED"
            diag = " [diagnostic]" if turn.is_diagnostic else ""
            print(
                f"  turn {turn.turn_idx}: action={turn.action} "
                f"motive={turn.motive_action} habit={turn.habit_action} "
                f"{tag}{diag} reward={turn.reward}"
            )
        print(f"  metrics: {trace.metrics}")
    return 0


def _cmd_compare(args: argparse.Namespace) -> int:
    import json

    from proteus.runtime.aggregate import aggregate_traces

    all_traces = []
    for path in args.trace_files:
        try:
            all_traces.extend(read_traces(path))
        except FileNotFoundError:
            print(f"trace file not found: {path}", file=sys.stderr)
            return 2
    if not all_traces:
        print("no traces to compare", file=sys.stderr)
        return 1

    groups = aggregate_traces(all_traces)
    for (model, difficulty), info in sorted(groups.items()):
        print(f"=== model={model} difficulty={difficulty} n={info['n']} ===")
        for key, mean in info["metrics"].items():
            print(f"  {key}: {mean:.2f}")

    if args.out:
        serializable = {
            f"{model}|{difficulty}": info
            for (model, difficulty), info in groups.items()
        }
        with open(args.out, "w") as fh:
            json.dump(serializable, fh, indent=2)
        print(f"summary written to {args.out}")
    return 0


def build_parser() -> argparse.ArgumentParser:
    """Build the top-level argument parser."""
    parser = argparse.ArgumentParser(
        prog="proteus",
        description="PROTEUS — a grid arena for measuring LLM motive-reading.",
    )
    sub = parser.add_subparsers(dest="command", required=True)

    run = sub.add_parser("run", help="run one session and append its trace")
    run.add_argument("--scenario", default="predator_evade")
    run.add_argument(
        "--model",
        required=True,
        help=(
            "provider spec '<name>:<model>'. Providers: "
            f"{', '.join(available_providers())}. Use 'fake:<name>' for offline."
        ),
    )
    run.add_argument("--difficulty", default="easy", choices=[d.value for d in Difficulty])
    run.add_argument("--seed", type=int, default=None)
    run.add_argument("--play-turns", type=int, default=15, dest="play_turns")
    run.add_argument("--no-probe", action="store_true", dest="no_probe")
    run.add_argument("--out", required=True, help="JSONL file to append the trace to")
    run.set_defaults(func=_cmd_run)

    play = sub.add_parser("play", help="play a session as a human via stdin")
    play.add_argument("--scenario", default="predator_evade")
    play.add_argument(
        "--difficulty", default="easy", choices=[d.value for d in Difficulty]
    )
    play.add_argument("--seed", type=int, default=None)
    play.add_argument("--play-turns", type=int, default=15, dest="play_turns")
    play.add_argument(
        "--probe",
        action="store_true",
        help="also ask the per-turn comprehension probe (default: off for humans)",
    )
    play.add_argument(
        "--out", default=None, help="optional JSONL file to append the human trace to"
    )
    play.set_defaults(func=_cmd_play)

    listing = sub.add_parser("list-scenarios", help="list registered scenarios")
    listing.set_defaults(func=_cmd_list_scenarios)

    replay = sub.add_parser("replay", help="print a saved trace")
    replay.add_argument("trace_file", help="path to a .jsonl trace file")
    replay.add_argument(
        "--visual", action="store_true", help="truecolor terminal replay"
    )
    replay.add_argument(
        "--png", default=None, metavar="DIR",
        help="also write per-frame PNGs to DIR (needs the 'viz' extra)",
    )
    replay.add_argument("--fps", type=float, default=4.0, help="replay frames/sec")
    replay.set_defaults(func=_cmd_replay)

    compare = sub.add_parser(
        "compare", help="aggregate traces by (model, difficulty) for baseline comparison"
    )
    compare.add_argument("trace_files", nargs="+", help="one or more .jsonl trace files")
    compare.add_argument(
        "--out", default=None, help="optional JSON file to write the aggregate summary to"
    )
    compare.set_defaults(func=_cmd_compare)

    return parser


def main(argv: list[str] | None = None) -> int:
    """CLI entry point. Returns a process exit code."""
    parser = build_parser()
    args = parser.parse_args(argv)
    return args.func(args)