proteus/cli/commands/compare.py

"""proteus.cli.commands.compare — the ``compare`` subcommand (aggregate traces)."""

from __future__ import annotations

import argparse
import sys

from proteus.game.runtime import read_traces


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

    from proteus.game.metrics.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