compare.py

"""Side-by-side agent comparison.

Runs multiple agents through the grader and prints a single tidy table
of results — overall score plus per-component metrics plus per-scenario
breakdown. Useful as the headline result of the project.

Usage::

    python compare.py
    python compare.py --agents random heuristic ppo --episodes 5
    python compare.py --output comparison.json
"""
from __future__ import annotations

import argparse
import json
import sys
from pathlib import Path
from typing import Callable

from agents import HeuristicAgent, RandomAgent
from grader import SCENARIOS, Grader


AgentFn = Callable[[dict], int]


def _load_agent(name: str) -> AgentFn:
    """Mirror :func:`evaluate._load_agent` so the two CLIs stay in sync."""
    if name == "random":
        return RandomAgent(seed=0)
    if name == "heuristic":
        return HeuristicAgent()
    if name == "ppo":
        # Defer to the same loader as evaluate.py to avoid drift.
        from evaluate import _load_agent as _ev_load_agent

        return _ev_load_agent("ppo")

    raise ValueError(f"Unknown agent: {name!r}")


# ---------------------------------------------------------------------
# Pretty printing
# ---------------------------------------------------------------------
def _bar(value: float, width: int = 12) -> str:
    """Return a unicode bar of length proportional to value (0-100)."""
    filled = int(round((max(0.0, min(100.0, value)) / 100.0) * width))
    return "█" * filled + "·" * (width - filled)


def _render_table(results: dict[str, dict]) -> str:
    """Render the side-by-side table as text."""
    lines: list[str] = []
    sep = "=" * 88
    sub = "-" * 88
    lines.append(sep)
    lines.append(" Hospital ED Resource Allocator — agent comparison")
    lines.append(sep)

    # ---- headline scores --------------------------------------------------
    header = f"  {'Agent':<12} {'Overall':>9} {'Survive':>8} {'Crit':>6} {'Wait':>7} {'Util':>6} {'Inval':>6}   bar"
    lines.append(header)
    lines.append(sub)
    for name, scores in results.items():
        bar = _bar(scores["overall_score"])
        lines.append(
            f"  {name:<12} "
            f"{scores['overall_score']:>9.2f} "
            f"{scores['survival_rate']*100:>7.1f}% "
            f"{scores['critical_patient_survival']*100:>5.1f}% "
            f"{scores['avg_wait_time']:>7.2f} "
            f"{scores['resource_utilization']*100:>5.1f}% "
            f"{scores['invalid_action_rate']*100:>5.1f}%   "
            f"{bar}"
        )
    lines.append(sub)

    # ---- per-scenario -----------------------------------------------------
    lines.append("  Per-scenario score:")
    scenarios = list(SCENARIOS.keys())
    header2 = "  " + " " * 12 + "".join(f"{s[:14]:>15}" for s in scenarios)
    lines.append(header2)
    for name, scores in results.items():
        row = f"  {name:<12}"
        for s in scenarios:
            row += f"{scores['scenario_scores'].get(s, float('nan')):>15.2f}"
        lines.append(row)
    lines.append(sep)
    return "\n".join(lines)


# ---------------------------------------------------------------------
# Main
# ---------------------------------------------------------------------
def _parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="Run multiple agents through the grader and print a comparison."
    )
    parser.add_argument(
        "--agents",
        nargs="+",
        default=["random", "heuristic", "ppo"],
        choices=["random", "heuristic", "ppo"],
        help="Which agents to compare. PPO is auto-skipped if no model exists.",
    )
    parser.add_argument(
        "--episodes",
        type=int,
        default=5,
        help="Episodes per scenario.",
    )
    parser.add_argument(
        "--seed",
        type=int,
        default=42,
        help="Base seed (deterministic).",
    )
    parser.add_argument(
        "--output",
        default=None,
        help="Optional JSON output file.",
    )
    return parser.parse_args()


def main() -> dict:
    args = _parse_args()
    grader = Grader(n_episodes_per_scenario=args.episodes, base_seed=args.seed)

    results: dict[str, dict] = {}
    for name in args.agents:
        agent = _load_agent(name)
        if agent is None:
            continue
        print(f"[compare] grading {name}…", file=sys.stderr)
        results[name] = grader.grade(agent)

    print(_render_table(results))

    if args.output:
        Path(args.output).write_text(json.dumps(results, indent=2))
        print(f"\n[compare] wrote {args.output}", file=sys.stderr)

    return results


if __name__ == "__main__":
    main()