"""Score the teacher, the base student and the distilled student on the same tasks.

Purpose: the measurement that decides whether the distillation was worth the hours.
    Three models, one task file, one set of sampling settings, run through Part 10's
    harness so the numbers are comparable with everything else in the course. Prints
    the gap that closed as a fraction of the gap that existed, and emits the rows in
    the shape the course's benchmark tables take, so a result can be pasted into a
    report without being retyped.
Platform: all (pure Python over HTTP; the three models are reached through an
    OpenAI-compatible API, so they may be served by any engine on any track, or by
    the Part 9 gateway one at a time under three aliases)
Minimum memory: 12 GB on the machine serving the largest of the three; this script
    needs very little
Assumes: Python 3.10 or newer; Part 10's run-eval.py and judge.py in --harness-dir;
    the three models reachable at --base-url under the names given, together or one
    at a time; distillog.py next to this file.

Usage: python3 evaluate-triplet.py --harness-dir ~/eval --tasks my-tasks.json \\
           --base-url http://127.0.0.1:4000/v1 \\
           --teacher local/teacher --base-student local/student-base \\
           --distilled local/student-distilled \\
           --quant Q4_K_M --engine llama.cpp --engine-version v0.4.0 \\
           --out-dir eval-out --labbook labbook.md

       Add --judge-model local/chat to grade all three runs with the same judge.
       Leave it out and only the deterministic checks are compared, which is a floor
       rather than a score but is exactly repeatable and needs no fourth model.

The three models must be served one at a time on a machine that cannot hold them
together. The script exits between models with a prompt so you can swap them, unless
--no-pause is given, which is what you want when a gateway loads them on demand.
"""

from __future__ import annotations

import argparse
import json
import subprocess
import sys
from pathlib import Path
from typing import Any

import distillog

ROLES = ("teacher", "base-student", "distilled-student")


def run(script: Path, arguments: list[str]) -> None:
    """Run one harness script, printing the command first so the log says what happened."""
    command = [sys.executable, str(script), *arguments]
    print("+ " + " ".join(command))
    subprocess.run(command, check=True)


def deterministic_rate(results_path: Path) -> tuple[int, int]:
    data = json.loads(results_path.read_text(encoding="utf-8"))
    passed = sum(1 for r in data["results"] if r["checks"]["passed"])
    return passed, len(data["results"])


def judged_mean(judged_path: Path) -> float | None:
    """Mean judge score, or None when a judge was not run or returned nothing usable.

    judge.py already computes this in its summary block; the per-result fallback
    exists so that a partially graded file still yields a number rather than an
    exception.
    """
    if not judged_path.is_file():
        return None
    data = json.loads(judged_path.read_text(encoding="utf-8"))
    summary = data.get("judge") or {}
    if isinstance(summary.get("judge_mean"), (int, float)):
        return summary["judge_mean"]
    scores = [r["judge"]["score"] for r in data.get("results", [])
              if isinstance(r.get("judge"), dict) and isinstance(r["judge"].get("score"), (int, float))]
    return round(sum(scores) / len(scores), 2) if scores else None


def per_category(results_path: Path) -> dict[str, tuple[int, int]]:
    data = json.loads(results_path.read_text(encoding="utf-8"))
    out: dict[str, list[int]] = {}
    for r in data["results"]:
        bucket = out.setdefault(r["category"], [0, 0])
        bucket[0] += int(r["checks"]["passed"])
        bucket[1] += 1
    return {k: (v[0], v[1]) for k, v in out.items()}


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__.splitlines()[0])
    parser.add_argument("--harness-dir", required=True, help="directory holding Part 10's run-eval.py and judge.py")
    parser.add_argument("--tasks", required=True, help="the Part 10 task file")
    parser.add_argument("--base-url", default="http://127.0.0.1:8080/v1")
    parser.add_argument("--api-key", default=None)
    parser.add_argument("--teacher", required=True, help="model name the endpoint serves the teacher under")
    parser.add_argument("--base-student", required=True, help="the student before distillation")
    parser.add_argument("--distilled", required=True, help="the student after distillation")
    parser.add_argument("--teacher-id", default=None, help="course model id of the teacher, for the record")
    parser.add_argument("--student-id", default=None, help="course model id of the student, for the record")
    parser.add_argument("--quant", default="unknown", help="quantisation actually loaded, for the record")
    parser.add_argument("--engine", default="llama.cpp")
    parser.add_argument("--engine-version", default="unknown")
    parser.add_argument("--judge-model", default=None, help="a fourth model that grades all three runs")
    parser.add_argument("--out-dir", default="eval-out")
    parser.add_argument("--no-pause", action="store_true",
                        help="do not stop between models; use when a gateway loads them on demand")
    parser.add_argument("--labbook", default=None)
    args = parser.parse_args()

    harness = Path(args.harness_dir)
    run_eval = harness / "run-eval.py"
    judge = harness / "judge.py"
    if not run_eval.is_file():
        raise SystemExit(f"{run_eval} not found; point --harness-dir at your Part 10 evaluation directory")
    if args.judge_model and not judge.is_file():
        raise SystemExit(f"{judge} not found, but --judge-model was given")

    out = Path(args.out_dir)
    out.mkdir(parents=True, exist_ok=True)
    models = {"teacher": args.teacher, "base-student": args.base_student, "distilled-student": args.distilled}

    summary: dict[str, dict[str, Any]] = {}
    for role in ROLES:
        name = models[role]
        if not args.no_pause and role != ROLES[0]:
            input(f"\nServe {name} now, then press Enter to evaluate it as the {role}. ")
        results = out / f"results-{role}.json"
        run(run_eval, [
            "--base-url", args.base_url,
            *(["--api-key", args.api_key] if args.api_key else []),
            "--model", name,
            "--quant", args.quant,
            "--engine", args.engine,
            "--engine-version", args.engine_version,
            "--tasks", args.tasks,
            "--out", str(results),
            "--notes", f"Part 15 triplet evaluation, {role}",
        ])
        judged = out / f"judged-{role}.json"
        if args.judge_model:
            run(judge, [
                "grade",
                "--base-url", args.base_url,
                *(["--api-key", args.api_key] if args.api_key else []),
                "--judge-model", args.judge_model,
                "--results", str(results),
                "--out", str(judged),
            ])
        passed, total = deterministic_rate(results)
        summary[role] = {
            "model": name,
            "checks_passed": passed,
            "checks_total": total,
            "judge_mean": judged_mean(judged),
            "by_category": per_category(results),
        }

    print("\n" + "=" * 72)
    print(f"{'role':<18}{'model':<26}{'checks':>10}{'judge mean':>14}")
    for role in ROLES:
        s = summary[role]
        judge_cell = "-" if s["judge_mean"] is None else f"{s['judge_mean']:.2f}"
        print(f"{role:<18}{s['model']:<26}{s['checks_passed']:>4}/{s['checks_total']:<5}{judge_cell:>14}")

    # The number the lab exists to produce: how much of the teacher-to-student gap
    # the distillation closed. Negative means the distilled student is worse than
    # the base it started from, which is a result and must be reported as one.
    base = summary["base-student"]["checks_passed"]
    tuned = summary["distilled-student"]["checks_passed"]
    teacher = summary["teacher"]["checks_passed"]
    gap = teacher - base
    closed = None if gap == 0 else round(100 * (tuned - base) / gap, 1)
    print()
    if gap == 0:
        print("The teacher and the base student passed the same number of checks: there was "
              "no gap on this task set, so there is nothing for distillation to close. "
              "Write a harder task set before drawing any conclusion.")
    else:
        print(f"gap teacher minus base student: {gap} check(s)")
        print(f"distilled student moved: {tuned - base} check(s)  ->  {closed}% of the gap closed")

    print("\nPer category, deterministic checks passed:")
    categories = sorted({c for role in ROLES for c in summary[role]["by_category"]})
    print(f"{'category':<14}{'teacher':>10}{'base':>10}{'distilled':>12}")
    for category in categories:
        cells = []
        for role in ROLES:
            got, of = summary[role]["by_category"].get(category, (0, 0))
            cells.append(f"{got}/{of}")
        print(f"{category:<14}{cells[0]:>10}{cells[1]:>10}{cells[2]:>12}")
    print("\nA category where the distilled student fell below the base student is the "
          "finding, whatever the total says. Part 15's challenge is about exactly that.")

    rows = [[role, summary[role]["model"],
             f"{summary[role]['checks_passed']}/{summary[role]['checks_total']}",
             summary[role]["judge_mean"] if summary[role]["judge_mean"] is not None else "not judged"]
            for role in ROLES]
    table_path = out / "benchmark-rows.json"
    table_path.write_text(json.dumps({
        "columns": ["Model", "Served as", "Deterministic checks", "Judge mean, 1-5"],
        "rows": rows,
        "context": {
            "engine": args.engine, "version": args.engine_version,
            "model": f"teacher {args.teacher_id or args.teacher}, student {args.student_id or args.base_student}",
            "quant": args.quant, "task_set": args.tasks,
        },
    }, indent=2), encoding="utf-8")
    print(f"\nbenchmark rows written to {table_path}")

    if args.labbook:
        rec = distillog.record(
            labbook=args.labbook,
            lab="part-15/evaluate-triplet",
            stage="evaluate",
            teacher={"id": args.teacher_id, "served_as": args.teacher},
            student={"id": args.student_id, "base_served_as": args.base_student,
                     "distilled_served_as": args.distilled},
            dataset={"tasks": args.tasks, "tasks_sha256": distillog.file_sha256(args.tasks)},
            hyperparameters={"engine": args.engine, "engine_version": args.engine_version,
                             "quant": args.quant, "judge_model": args.judge_model},
            seed=None,
            scores={
                "teacher": summary["teacher"],
                "base_student": summary["base-student"],
                "distilled_student": summary["distilled-student"],
                "gap": gap,
                "percent_of_gap_closed": closed,
            },
            config_path=__file__,
            notes=None,
        )
        print(f"recorded {rec['stage']} as {rec['run_id']} in {args.labbook}")


if __name__ == "__main__":
    main()