#!/usr/bin/env python3
"""Assemble the per-quantisation results into the course comparison table, and pick a file.

Purpose: read everything measure-quants.sh wrote - the divergence summaries, the llama-bench
    output and the Part 10 task-set results - into one table with one row per quantisation, then
    apply the two thresholds you set (a divergence budget and a memory budget) and name the
    fastest file that meets both. The recommendation is arithmetic on your own numbers, not an
    opinion: change the thresholds and it changes.
Platform: all (pure Python, standard library only)
Minimum memory: negligible
Assumes: Python 3.9 or later and a results directory written by measure-quants.sh containing
    kld-*.json files, and optionally bench-*.json and eval-*.json for the speed and task columns.

Usage: python3 measure-quants.py --results results --labbook labbook.md
       python3 measure-quants.py --results results --quant-dir ~/models/gguf/five-ways \
           --max-kld 0.02 --budget-gb 8 --labbook labbook.md
"""

from __future__ import annotations

import argparse
import json
import platform
import time
from pathlib import Path
from typing import Optional

DASH = "—"


def load(path: Path) -> Optional[dict]:
    try:
        return json.loads(path.read_text(encoding="utf-8"))
    except (OSError, ValueError):
        return None


def bench_rates(payload) -> tuple:
    """Pull prefill and decode rates out of llama-bench's JSON, without assuming field names.

    llama-bench reports one record per test. A record with n_gen == 0 is the prefill test and one
    with n_prompt == 0 is the generation test; the throughput field is avg_ts in the versions this
    course was written against. Anything unrecognised comes back as None rather than as a guess.
    """
    if not isinstance(payload, list):
        return None, None
    prefill = decode = None
    for record in payload:
        if not isinstance(record, dict):
            continue
        rate = record.get("avg_ts")
        if rate is None:
            rate = next((v for k, v in record.items()
                         if k.endswith("_ts") and isinstance(v, (int, float))), None)
        if rate is None:
            continue
        if record.get("n_gen") in (0, None) and record.get("n_prompt"):
            prefill = float(rate)
        elif record.get("n_gen"):
            decode = float(rate)
    return prefill, decode


def fmt(value, digits: int = 3) -> str:
    if value is None:
        return DASH
    if isinstance(value, float):
        return f"{value:.{digits}f}"
    return str(value)


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__.splitlines()[0])
    parser.add_argument("--results", default="results", help="directory written by measure-quants.sh")
    parser.add_argument("--quant-dir", default=None, help="directory of GGUF files, for file sizes")
    parser.add_argument("--reference", default=None,
                        help="the full-precision GGUF, so the reference row carries its size")
    parser.add_argument("--max-kld", type=float, default=None,
                        help="the largest mean divergence you are willing to accept; rows above it "
                             "are excluded from the recommendation")
    parser.add_argument("--min-same-top", type=float, default=None,
                        help="the smallest acceptable share of positions where the quantised model "
                             "still puts the same token on top")
    parser.add_argument("--budget-gb", type=float, default=None,
                        help="memory you can spend on weights; rows larger than this are excluded")
    parser.add_argument("--out", default=None, help="write the assembled table as JSON")
    parser.add_argument("--labbook", default=None)
    args = parser.parse_args()

    results = Path(args.results)
    if not results.is_dir():
        raise SystemExit(f"{results} is not a directory; run measure-quants.sh first")

    sizes = {}
    if args.quant_dir:
        for path in Path(args.quant_dir).glob("*.gguf"):
            sizes[path.stem] = path.stat().st_size / 1e9

    reference_gb = None
    if args.reference:
        reference_path = Path(args.reference)
        if reference_path.is_file():
            reference_gb = round(reference_path.stat().st_size / 1e9, 2)

    reference_checks = reference_tasks = None
    reference_eval = load(results / "eval-BF16.json")
    if reference_eval and "run" in reference_eval:
        reference_checks = reference_eval["run"].get("checks_passed")
        reference_tasks = reference_eval["run"].get("tasks")

    rows = []
    for kld_file in sorted(results.glob("kld-*.json")):
        label = kld_file.stem[len("kld-"):]
        payload = load(kld_file) or {}
        summary = payload.get("summary", {})

        prefill, decode = bench_rates(load(results / f"bench-{label}.json"))

        checks = judged = tasks = None
        evaluation = load(results / f"eval-{label}.json")
        if evaluation and "run" in evaluation:
            checks = evaluation["run"].get("checks_passed")
            tasks = evaluation["run"].get("tasks")
        judged_file = load(results / f"judged-{label}.json")
        if judged_file and "judge" in judged_file:
            judged = judged_file["judge"].get("judge_mean")

        rows.append({
            "quantisation": label,
            "file_gb": round(sizes[label], 2) if label in sizes else None,
            "mean_kld": summary.get("mean_kld"),
            "p99_kld": summary.get("p99_kld"),
            "same_top_token": summary.get("same_top_token"),
            "checks_passed": checks,
            "tasks": tasks,
            "judge_mean": judged,
            "prefill_tokens_per_s": prefill,
            "decode_tokens_per_s": decode,
        })

    if not rows:
        raise SystemExit(f"no kld-*.json files in {results}; measure-quants.sh writes them")

    header = ["Quantisation", "File GB", "Mean KLD", "99% KLD", "Same top",
              "Checks", "Judge", "pp512 t/s", "tg128 t/s"]
    print("\n| " + " | ".join(header) + " |")
    print("| " + " | ".join("---" for _ in header) + " |")
    reference_column = (f"{reference_checks}/{reference_tasks}"
                        if reference_checks is not None and reference_tasks else DASH)
    print(f"| BF16 (reference) | {fmt(reference_gb, 2)} | 0 | 0 | 1.0000 "
          f"| {reference_column} | {DASH} | {DASH} | {DASH} |")
    for row in sorted(rows, key=lambda r: (r["mean_kld"] is None, r["mean_kld"] or 0)):
        checks = (f"{row['checks_passed']}/{row['tasks']}"
                  if row["checks_passed"] is not None and row["tasks"] else DASH)
        print("| " + " | ".join([
            row["quantisation"],
            fmt(row["file_gb"], 2),
            fmt(row["mean_kld"], 5),
            fmt(row["p99_kld"], 5),
            fmt(row["same_top_token"], 4),
            checks,
            fmt(row["judge_mean"], 2),
            fmt(row["prefill_tokens_per_s"], 1),
            fmt(row["decode_tokens_per_s"], 1),
        ]) + " |")

    # ---- the choice ------------------------------------------------------
    eligible = list(rows)
    reasons = []
    if args.max_kld is not None:
        eligible = [r for r in eligible if r["mean_kld"] is not None and r["mean_kld"] <= args.max_kld]
        reasons.append(f"mean KLD at most {args.max_kld}")
    if args.min_same_top is not None:
        eligible = [r for r in eligible
                    if r["same_top_token"] is not None and r["same_top_token"] >= args.min_same_top]
        reasons.append(f"same top token at least {args.min_same_top}")
    if args.budget_gb is not None:
        eligible = [r for r in eligible if r["file_gb"] is not None and r["file_gb"] <= args.budget_gb]
        reasons.append(f"file at most {args.budget_gb} GB")

    choice = None
    if reasons:
        print(f"\nThresholds applied: {'; '.join(reasons)}.")
        with_speed = [r for r in eligible if r["decode_tokens_per_s"] is not None]
        if with_speed:
            choice = max(with_speed, key=lambda r: r["decode_tokens_per_s"])
            print(f"Fastest file meeting all of them: {choice['quantisation']}.")
        elif eligible:
            choice = min(eligible, key=lambda r: r["mean_kld"] or 0)
            print(f"No speed measurements; least divergent file meeting the thresholds: "
                  f"{choice['quantisation']}.")
        else:
            print("No file meets all the thresholds. Either the budget is too tight or the "
                  "quantisations are worse than you assumed; both are results worth recording.")
    else:
        print("\nNo thresholds given, so no recommendation. Pass --max-kld, --min-same-top and "
              "--budget-gb to turn the table into a decision.")

    record = {
        "lab": "part-16/lab-quantise-five-ways-and-measure/table",
        "run_id": time.strftime("%Y%m%dT%H%M%S"),
        "reference": {"file_gb": reference_gb, "checks_passed": reference_checks,
                      "tasks": reference_tasks},
        "rows": rows,
        "thresholds": {"max_kld": args.max_kld, "min_same_top": args.min_same_top,
                       "budget_gb": args.budget_gb},
        "choice": choice["quantisation"] if choice else None,
        "host": platform.platform(),
        "date": time.strftime("%Y-%m-%d"),
    }

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

    print("\nBefore you act on the ordering: run one quantisation twice and compare the two "
          "\nrows. Any difference smaller than that gap is not a difference you have measured.")


if __name__ == "__main__":
    main()