#!/usr/bin/env python3
"""Run the same task suite across scaffolds and models, and report which one moved the result.

Purpose: the reality check's instrument. It calls agent-eval.py once per (scaffold, model)
    cell, collects the result files, and prints the grid plus the two contrasts the claim
    "agents are just loops" turns on: how much the result moves when the scaffold is fixed
    and the model changes, and how much it moves when the model is fixed and the scaffold
    changes. It computes nothing the underlying runs did not measure.
Platform: all (pure Python; the models may be served on any track or machine)
Minimum memory: 16 GB on the machine serving the largest model; this script needs almost none
Assumes: Python 3.9 or later, agent-eval.py and the scaffold entry points in the same
    directory, and every model alias already answering on --base-url. Each cell is a full
    suite run, so the wall clock is (scaffolds x models x tasks x repeats) agent runs:
    check the estimate this script prints before you start it.

Usage: python3 scaffold-comparison.py --base-url http://127.0.0.1:4000/v1 \\
           --scaffold minimal=scaffold-minimal.py \\
           --scaffold four-role=multi-agent-system.py \\
           --scaffold smolagents=scaffold-smolagents.py \\
           --model local/small --model local/mid --model local/answer \\
           --workspace ./agent-workspace --out comparison.json --labbook labbook.md
"""

from __future__ import annotations

import argparse
import json
import statistics
import subprocess
import sys
import time
from pathlib import Path
from typing import Dict, List, Optional


def spread(values: List[float]) -> Optional[float]:
    """Max minus min. The plainest measure of "how much did this factor move the result"."""
    return round(max(values) - min(values), 3) if values else None


def run_cell(args, scaffold_name: str, entry_point: str, model: str) -> Optional[dict]:
    """One full suite run, as a subprocess, so a crash in one cell does not end the grid."""
    safe = f"{scaffold_name}-{model.replace('/', '_')}"
    out = Path(args.results_dir) / f"result-{safe}.json"
    if out.exists() and args.skip_existing:
        print(f"  reusing {out}")
        return json.loads(out.read_text(encoding="utf-8"))

    command = [sys.executable, str(Path(args.eval_script)),
               "--agent", entry_point,
               "--tasks", args.tasks,
               "--model", model,
               "--base-url", args.base_url,
               "--repeats", str(args.repeats),
               "--out", str(out),
               "--trajectory-dir", args.trajectory_dir,
               "--notes", f"scaffold comparison cell {scaffold_name} x {model}"]
    if args.router_model:
        command += ["--router-model", args.router_model]
    if args.workspace:
        command += ["--workspace", args.workspace]
    if args.index:
        command += ["--index", args.index]
    if args.only:
        command += ["--only", args.only]
    if args.cost_per_million is not None:
        command += ["--cost-per-million", str(args.cost_per_million)]
    for option in args.option:
        command += ["--option", option]

    print(f"  {' '.join(command[1:])}")
    finished = subprocess.run(command, check=False)
    if not out.exists():
        print(f"  no result file for {scaffold_name} x {model} "
              f"(exit status {finished.returncode}); recorded as a missing cell")
        return None
    return json.loads(out.read_text(encoding="utf-8"))


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__.splitlines()[0])
    parser.add_argument("--scaffold", action="append", default=[], metavar="NAME=FILE",
                        help="a scaffold to test; repeat for each one")
    parser.add_argument("--model", action="append", default=[], metavar="ALIAS",
                        help="a model alias to test; repeat for each one")
    parser.add_argument("--router-model", default=None,
                        help="alias of the small router, for scaffolds that use one")
    parser.add_argument("--tasks", default="agent-tasks.json")
    parser.add_argument("--only", default=None, help="run one task category only")
    parser.add_argument("--repeats", type=int, default=1,
                        help="attempts per task per cell; 3 makes the variance visible")
    parser.add_argument("--base-url", default="http://127.0.0.1:4000/v1")
    parser.add_argument("--workspace", default=None)
    parser.add_argument("--index", default=None)
    parser.add_argument("--option", action="append", default=[], metavar="KEY=VALUE",
                        help="passed through to every agent entry point")
    parser.add_argument("--cost-per-million", type=float, default=None,
                        help="your figure from Part 23's cost-model.py; no default")
    parser.add_argument("--eval-script", default="agent-eval.py")
    parser.add_argument("--results-dir", default="comparison-results")
    parser.add_argument("--trajectory-dir", default="trajectories")
    parser.add_argument("--skip-existing", action="store_true",
                        help="reuse a cell's result file if it is already there")
    parser.add_argument("--out", default="comparison.json")
    parser.add_argument("--labbook", default=None)
    args = parser.parse_args()

    if not args.scaffold or not args.model:
        parser.error("give at least one --scaffold NAME=FILE and one --model ALIAS")

    scaffolds = []
    for item in args.scaffold:
        if "=" not in item:
            parser.error(f"--scaffold needs NAME=FILE, got {item!r}")
        name, path = item.split("=", 1)
        scaffolds.append((name.strip(), path.strip()))

    spec = json.loads(Path(args.tasks).read_text(encoding="utf-8"))
    task_count = len([t for t in spec["tasks"]
                      if not args.only or t.get("category") == args.only])
    cells = len(scaffolds) * len(args.model)
    print(f"{cells} cell(s): {len(scaffolds)} scaffold(s) x {len(args.model)} model(s)")
    print(f"{task_count} task(s) x {args.repeats} attempt(s) per cell = "
          f"{cells * task_count * args.repeats} agent run(s) in total\n")

    Path(args.results_dir).mkdir(parents=True, exist_ok=True)
    started = time.time()
    grid: Dict[str, Dict[str, dict]] = {}

    for name, path in scaffolds:
        grid[name] = {}
        for model in args.model:
            print(f"{name} x {model}")
            payload = run_cell(args, name, path, model)
            if payload is None:
                grid[name][model] = {"missing": True}
                continue
            summary = payload["run"]["summary"]
            grid[name][model] = {
                "success_rate": summary["success_rate"],
                "mean_steps": summary["mean_steps"],
                "mean_tokens": summary["mean_tokens"],
                "mean_seconds": summary["mean_seconds"],
                "total_tokens": summary["total_tokens"],
                "cost_per_task": summary.get("cost_per_task"),
                "result_file": str(Path(args.results_dir) / f"result-{name}-"
                                   f"{model.replace('/', '_')}.json"),
            }
            print("")

    # The two contrasts. Each is the spread of success rate across one axis with the
    # other held fixed, which is exactly what the claim under test is about.
    by_scaffold = {}
    for name in grid:
        rates = [c["success_rate"] for c in grid[name].values() if "success_rate" in c]
        by_scaffold[name] = {"model_spread": spread(rates),
                             "mean_success": round(statistics.fmean(rates), 3) if rates else None}
    by_model = {}
    for model in args.model:
        rates = [grid[n][model]["success_rate"] for n in grid
                 if "success_rate" in grid[n].get(model, {})]
        by_model[model] = {"scaffold_spread": spread(rates),
                           "mean_success": round(statistics.fmean(rates), 3) if rates else None}

    report = {
        "lab": "part-26/reality-check-agents-are-just-loops",
        "run_id": time.strftime("%Y%m%dT%H%M%S"),
        "date": time.strftime("%Y-%m-%d"),
        "task_set": spec.get("name", args.tasks),
        "task_set_version": spec.get("version"),
        "tasks": task_count,
        "repeats": args.repeats,
        "scaffolds": [n for n, _ in scaffolds],
        "models": args.model,
        "grid": grid,
        "changing_the_model": by_scaffold,
        "changing_the_scaffold": by_model,
        "wall_seconds": round(time.time() - started, 1),
    }
    Path(args.out).write_text(json.dumps(report, indent=2), encoding="utf-8")

    width = max(len(n) for n, _ in scaffolds) + 2
    print("\nsuccess rate")
    print(" " * width + "".join(f"{m:>22}" for m in args.model))
    for name, _ in scaffolds:
        row = "".join(f"{grid[name][m].get('success_rate', float('nan')):>22.2f}"
                      if "success_rate" in grid[name].get(m, {}) else f"{'missing':>22}"
                      for m in args.model)
        print(f"{name:<{width}}{row}")

    print("\nmean tokens per task")
    print(" " * width + "".join(f"{m:>22}" for m in args.model))
    for name, _ in scaffolds:
        row = "".join(f"{grid[name][m].get('mean_tokens', 0):>22.0f}"
                      if "mean_tokens" in grid[name].get(m, {}) else f"{'missing':>22}"
                      for m in args.model)
        print(f"{name:<{width}}{row}")

    print("\nhow much each factor moved the success rate (max minus min)")
    for name in by_scaffold:
        print(f"  changing the model, scaffold fixed at {name:<16} "
              f"{by_scaffold[name]['model_spread']}")
    for model in by_model:
        print(f"  changing the scaffold, model fixed at {model:<16} "
              f"{by_model[model]['scaffold_spread']}")
    print(f"\nwritten to {args.out}")

    if args.labbook:
        with Path(args.labbook).open("a", encoding="utf-8") as handle:
            handle.write(json.dumps(report) + "\n")
        print(f"recorded in {args.labbook}")


if __name__ == "__main__":
    main()