scripts/benchmark_medflow.py

"""Benchmark MedFlow-OpenEnv with a few lightweight policies.

Usage:
    python scripts/benchmark_medflow.py --episodes 3

Outputs:
    - outputs/benchmarks/benchmark_summary.json
    - outputs/benchmarks/benchmark_summary.md

The benchmark compares a greedy policy against a random policy across all
defined tasks and records the environment's final score and episode stats.
"""
from __future__ import annotations

import argparse
import json
import random
import sys
from pathlib import Path
from statistics import mean
from typing import Dict, Iterable, List

ROOT = Path(__file__).resolve().parents[1]
if str(ROOT) not in sys.path:
    sys.path.insert(0, str(ROOT))

from app.env import HospitalQueueEnvironment
from app.models import HospitalAction
from app.tasks import ALL_TASKS


def _choose_greedy_action(obs) -> HospitalAction:
    waiting = list(obs.waiting_patients)
    doctors = list(obs.doctors)
    free_docs = [d for d in doctors if not d.get("busy", False)]

    if waiting and free_docs:
        priority_rank = {"emergency": 0, "urgent": 1, "normal": 2}
        patient = sorted(
            waiting,
            key=lambda p: (
                priority_rank.get(p.get("priority", "normal"), 3),
                -int(p.get("severity_score", 0)),
                -int(p.get("wait_minutes", 0)),
                int(p.get("id", 0)),
            ),
        )[0]
        required = patient.get("required_specialization")
        matching = [d for d in free_docs if d.get("specialization") == required]
        general = [d for d in free_docs if d.get("specialization") == "General"]
        doctor = (matching or general or free_docs)[0]
        return HospitalAction(action_type="assign", patient_id=int(patient["id"]), doctor_id=int(doctor["id"]))

    if waiting:
        patient = waiting[0]
        return HospitalAction(action_type="prioritize", patient_id=int(patient["id"]))

    for d in doctors:
        if (not d.get("busy", False)) and d.get("current_patient_id") is not None:
            return HospitalAction(action_type="discharge", patient_id=int(d["current_patient_id"]))

    return HospitalAction(action_type="wait")


def _choose_random_action(obs) -> HospitalAction:
    waiting = list(obs.waiting_patients)
    doctors = list(obs.doctors)
    free_docs = [d for d in doctors if not d.get("busy", False)]
    if waiting and free_docs:
        patient = random.choice(waiting)
        doctor = random.choice(free_docs)
        return HospitalAction(action_type="assign", patient_id=int(patient["id"]), doctor_id=int(doctor["id"]))
    if waiting:
        return HospitalAction(action_type="prioritize", patient_id=int(random.choice(waiting)["id"]))
    for d in doctors:
        if d.get("current_patient_id") is not None:
            return HospitalAction(action_type="discharge", patient_id=int(d["current_patient_id"]))
    return HospitalAction(action_type="wait")


def run_episode(task_id: str, policy: str, seed: int = 42) -> Dict[str, object]:
    random.seed(seed)
    env = HospitalQueueEnvironment()
    obs = env.reset(task_id=task_id, seed=seed)
    total_reward = 0.0
    steps = 0

    while not obs.done and steps < ALL_TASKS[task_id].max_steps:
        if policy == "greedy":
            action = _choose_greedy_action(obs)
        elif policy == "random":
            action = _choose_random_action(obs)
        else:
            raise ValueError(f"Unknown policy: {policy}")
        obs = env.step(action)
        total_reward += float(obs.reward or 0.0)
        steps += 1

    state = env.state
    return {
        "task_id": task_id,
        "policy": policy,
        "final_score": float(state.final_score or 0.0),
        "total_reward": round(total_reward, 3),
        "steps": steps,
        "total_seen": state.total_patients_seen,
        "total_arrived": state.total_patients_arrived,
        "avg_wait_all": float(state.avg_wait_all),
        "critical_deaths": len(state.emergency_response_times),
        "submitted": bool(state.submitted),
    }


def _write_markdown(rows: List[Dict[str, object]], out_path: Path) -> None:
    lines = [
        "# MedFlow OpenEnv Benchmark Summary",
        "",
        "| task_id | policy | final_score | total_reward | steps | seen/arrived | avg_wait_all |",
        "| --- | --- | ---: | ---: | ---: | ---: | ---: |",
    ]
    for row in rows:
        lines.append(
            f"| {row['task_id']} | {row['policy']} | {row['final_score']:.3f} | {row['total_reward']:.3f} | {row['steps']} | {row['total_seen']}/{row['total_arrived']} | {row['avg_wait_all']:.2f} |"
        )
    out_path.write_text("\n".join(lines), encoding="utf-8")


def _summarize(rows: List[Dict[str, object]]) -> Dict[str, object]:
    by_policy: Dict[str, List[Dict[str, object]]] = {}
    for row in rows:
        by_policy.setdefault(str(row["policy"]), []).append(row)

    summary: Dict[str, object] = {"policies": {}, "rows": rows}
    for policy, policy_rows in by_policy.items():
        summary["policies"][policy] = {
            "avg_final_score": round(mean(float(r["final_score"]) for r in policy_rows), 3),
            "avg_total_reward": round(mean(float(r["total_reward"]) for r in policy_rows), 3),
            "avg_steps": round(mean(int(r["steps"]) for r in policy_rows), 1),
            "success_rate": round(sum(1 for r in policy_rows if float(r["final_score"]) >= 0.10) / len(policy_rows), 3),
        }
    return summary


def main(episodes: int, outdir: Path, seed: int) -> None:
    outdir.mkdir(parents=True, exist_ok=True)
    random.seed(seed)

    rows: List[Dict[str, object]] = []
    for task_id in ALL_TASKS.keys():
        for policy in ("greedy", "random"):
            for episode in range(episodes):
                row = run_episode(task_id=task_id, policy=policy, seed=seed + episode)
                rows.append(row)
                print(
                    f"{task_id} | {policy} | score={row['final_score']:.3f} | reward={row['total_reward']:.3f} | steps={row['steps']}"
                )

    summary = _summarize(rows)
    summary_path = outdir / "benchmark_summary.json"
    markdown_path = outdir / "benchmark_summary.md"
    summary_path.write_text(json.dumps(summary, indent=2), encoding="utf-8")
    _write_markdown(rows, markdown_path)

    print(f"Saved JSON summary to {summary_path}")
    print(f"Saved markdown summary to {markdown_path}")
    for policy, stats in summary["policies"].items():
        print(
            f"{policy}: avg_final_score={stats['avg_final_score']:.3f}, avg_total_reward={stats['avg_total_reward']:.3f}, success_rate={stats['success_rate']:.3f}"
        )


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--episodes", type=int, default=3)
    parser.add_argument("--seed", type=int, default=42)
    parser.add_argument("--outdir", type=Path, default=Path("outputs/benchmarks"))
    args = parser.parse_args()
    main(episodes=args.episodes, outdir=args.outdir, seed=args.seed)