#!/usr/bin/env python3
"""Summarise this lab's llama-bench records against the Part 5 decode prediction.

Purpose: read the notebook lines bench-to-labbook.py appended, keep the latest measurement
    of each model file on each backend, and print per model: the prompt-processing and
    generation rates with their spread, the ratio between the two, the decode ceiling Part 5's
    arithmetic gives for the exact file measured, the share of that ceiling reached, and the
    time per token the ceiling does not explain, and the arithmetic rate prefill implies.
    Generation measured at a depth is set beside the slowdown the KV cache formula
    predicts. With --format benchmark the same rows are printed as a <Benchmark> block,
    the course benchmark format.
Platform: all (pure Python standard library; nothing platform-specific)
Minimum memory: 8 GB
Assumes: Python 3.9 or later; labbook.md holding lines written by bench-to-labbook.py; the
    copy of the course's models.json that Part 5's predict-decode.py read; a bandwidth figure,
    passed with --bandwidth-gbps or taken from the latest part-05/predict-decode line.

Usage: python3 summarise-benchmarks.py --labbook labbook.md --models models.json
       python3 summarise-benchmarks.py --labbook labbook.md --models models.json \
           --format benchmark --hardware "Track S, DGX Spark, 128 GB" --os "DGX OS 7"

Rows are grouped by backend and file. A record made with GGML_CUDA_ENABLE_UNIFIED_MEMORY set
(bench-to-labbook.py writes the field) is shown as backend "+UMA", so a HIP run with the
variable never replaces the same file's run without it.

Method: ceiling = bandwidth / active bytes per token, where active bytes = the size
llama-bench reports for the file x active parameters / total parameters (the shortcut
predict-decode.py uses; 1 GB = 1e9 bytes). The unexplained time per token is
1 / measured rate - 1 / ceiling. Prefill TFLOP/s = prompt tokens per second x 2 x active
parameters. Depth prediction: rate(t) / rate(0) = W / (W + k x t), with W the active bytes
and k the KV bytes per token from models.json, scaled for the cache type.
"""

from __future__ import annotations

import argparse
import json
import sys
from pathlib import Path

LAB = "part-06/lab-benchmark-the-reference-models"
# Bytes per cached element, from the ggml block layouts: f16 is 2 bytes; q8_0 stores 32
# values in 34 bytes; q4_0 stores 32 in 18.
BYTES_PER_ELEMENT = {"f32": 4.0, "f16": 2.0, "bf16": 2.0, "q8_0": 34 / 32, "q4_0": 18 / 32}


def read_records(path: Path) -> list[dict]:
    records = []
    for line in path.read_text(encoding="utf-8").splitlines():
        line = line.strip()
        if not line.startswith("{"):
            continue  # prose, headings and tables in the notebook
        try:
            obj = json.loads(line)
        except json.JSONDecodeError:
            continue
        if isinstance(obj, dict):
            records.append(obj)
    return records


def match_model(catalogue: list[dict], model_path: str) -> dict | None:
    """The catalogue entry whose name starts the file name, longest name first."""
    base = Path(model_path or "").name.lower()
    hits = [m for m in catalogue if base.startswith(str(m.get("name", "")).lower() + "-")]
    return max(hits, key=lambda m: len(m["name"])) if hits else None


def backend_label(record: dict):
    """The backend string, marked when llama.cpp ran with GGML_CUDA_ENABLE_UNIFIED_MEMORY set."""
    backend = record.get("backend")
    if record.get("ggml_cuda_enable_unified_memory") is not None:
        return f"{backend}+UMA"
    return backend


def fmt(value, digits=2, width=8) -> str:
    return f"{value:>{width}.{digits}f}" if isinstance(value, (int, float)) else f"{'-':>{width}}"


def main() -> int:
    parser = argparse.ArgumentParser(description=__doc__.splitlines()[0])
    parser.add_argument("--labbook", required=True, help="the notebook bench-to-labbook.py wrote to")
    parser.add_argument("--models", help="the course models.json (needed for ceilings)")
    parser.add_argument("--bandwidth-gbps", type=float, help="measured bandwidth in GB/s (Part 5)")
    parser.add_argument("--format", choices=["table", "benchmark"], default="table")
    parser.add_argument("--hardware", default="FILL-IN: track, machine, memory")
    parser.add_argument("--os", default="FILL-IN: operating system and version")
    parser.add_argument("--record", action="store_true",
                        help="append one part-06/compare-with-prediction line to the notebook")
    args = parser.parse_args()

    notebook = Path(args.labbook)
    if not notebook.exists():
        print(f"{notebook} does not exist", file=sys.stderr)
        return 1
    records = read_records(notebook)
    bench = [r for r in records if r.get("lab") == LAB and isinstance(r.get("tokens_per_s"), (int, float))]
    if not bench:
        print(f"no {LAB} lines with a rate in {notebook}; run the sweep first", file=sys.stderr)
        return 1

    bandwidth, source = args.bandwidth_gbps, "--bandwidth-gbps"
    if bandwidth is None:
        for r in records:
            if r.get("lab") == "part-05/predict-decode" and r.get("bandwidth_gbps"):
                bandwidth, source = float(r["bandwidth_gbps"]), "the latest part-05/predict-decode line"
    catalogue = []
    if args.models:
        catalogue = json.loads(Path(args.models).read_text(encoding="utf-8")).get("models", [])

    # A partial-offload experiment on the same file is a different measurement: keep only the
    # records made with the largest layer count asked for on each file and backend.
    most_layers: dict[tuple, int] = {}
    for r in bench:
        key = (backend_label(r), Path(str(r.get("model_path") or r.get("model"))).name)
        most_layers[key] = max(most_layers.get(key, -1), int(r.get("n_gpu_layers") or 0))
    bench = [r for r in bench if int(r.get("n_gpu_layers") or 0) == most_layers[
        (backend_label(r), Path(str(r.get("model_path") or r.get("model"))).name)]]

    # Latest record per (backend, file, test); later lines in the notebook win.
    latest: dict[tuple, dict] = {}
    for r in bench:
        key = (backend_label(r), Path(str(r.get("model_path") or r.get("model"))).name, r.get("test"))
        latest[key] = r
    groups: dict[tuple, dict] = {}
    for (backend, name, test), r in latest.items():
        groups.setdefault((backend, name), {})[test] = r

    rows, depth_rows = [], []
    for (backend, name), tests in sorted(groups.items(), key=lambda kv: (str(kv[0][0]), kv[0][1])):
        # At depth 0, the most recently measured prompt-processing and generation tests.
        by_time = sorted(tests.items(), key=lambda kv: str(kv[1].get("measured_on") or ""))
        pp = next((r for t, r in reversed(by_time) if t.startswith("pp") and "@" not in t and "+" not in t), None)
        tg = next((r for t, r in reversed(by_time) if t.startswith("tg") and "@" not in t), None)
        any_row = tg or pp or next(iter(tests.values()))
        entry = match_model(catalogue, any_row.get("model_path") or name)
        size_bytes = any_row.get("model_size_bytes") or (any_row.get("file_gb") or 0) * 1e9
        active_gb = ceiling = unexplained = reached = prefill_tflops = None
        if entry and size_bytes:
            params = entry.get("params", {})
            if pp and params.get("activeB"):
                # Two floating-point operations per active parameter per token, as in Part 3.
                prefill_tflops = pp["tokens_per_s"] * 2 * float(params["activeB"]) * 1e9 / 1e12
            fraction = min(float(params.get("activeB", 0)) / float(params.get("totalB", 1)), 1.0) or 1.0
            active_gb = size_bytes * fraction / 1e9
            if bandwidth:
                ceiling = bandwidth / active_gb
        rate = tg["tokens_per_s"] if tg else None
        if rate and ceiling:
            reached = rate / ceiling
            unexplained = 1000.0 / rate - 1000.0 / ceiling
        spread = tg.get("tokens_per_s_stddev") if tg else None
        rows.append({
            "backend": backend, "file": name, "id": entry.get("id") if entry else None,
            "quant": any_row.get("quant"), "file_gb": round(size_bytes / 1e9, 2) if size_bytes else None,
            "active_gb": round(active_gb, 2) if active_gb else None,
            "pp_test": pp.get("test") if pp else None, "pp": pp["tokens_per_s"] if pp else None,
            "tg_test": tg.get("test") if tg else None, "tg": rate, "tg_stddev": spread,
            "tg_cv_pct": 100.0 * spread / rate if rate and isinstance(spread, (int, float)) else None,
            "pp_over_tg": pp["tokens_per_s"] / rate if pp and rate else None,
            "ceiling": ceiling, "reached": reached, "unexplained_ms": unexplained,
            "prefill_tflops": prefill_tflops,
            "build": any_row.get("build"), "build_number": any_row.get("build_number"),
            "measured_on": any_row.get("measured_on"),
            "context": max(int(r.get("n_prompt") or 0) + int(r.get("n_gen") or 0) + int(r.get("n_depth") or 0)
                           for r in tests.values()),
        })
        for test, r in sorted(tests.items(), key=lambda kv: int(kv[1].get("n_depth") or 0)):
            depth = int(r.get("n_depth") or 0)
            base = tests.get(f"tg{r.get('n_gen')}")  # the same generation test at depth 0
            if not test.startswith("tg") or depth == 0 or not base:
                continue
            predicted = None
            if entry and active_gb and entry.get("kv", {}).get("bytesPerTokenFp16"):
                per_element = BYTES_PER_ELEMENT.get(str(r.get("type_k") or "f16"), 2.0)
                k = entry["kv"]["bytesPerTokenFp16"] / 2.0 * per_element
                predicted = active_gb * 1e9 / (active_gb * 1e9 + k * depth)
            depth_rows.append({"file": name, "backend": backend, "test": test, "rate": r["tokens_per_s"],
                               "measured_ratio": r["tokens_per_s"] / base["tokens_per_s"], "predicted_ratio": predicted,
                               "architecture": entry.get("architecture") if entry else None})

    if args.format == "benchmark":
        context = max(r["context"] for r in rows)
        builds = sorted({f"build {r['build_number']} ({r['build']})" for r in rows})
        dates = sorted(str(r["measured_on"])[:10] for r in rows if r["measured_on"])
        print("<Benchmark")
        print('  title="llama-bench, prompt processing and generation, one row per model file"')
        print("  columns={['Model file', 'Backend', 'File GB', 'pp tokens/s', 'tg tokens/s', "
              "'tg ceiling tokens/s', 'Share of ceiling']}")
        print("  rows={[")
        for r in rows:
            cells = [r["file"], r["backend"], r["file_gb"],
                     round(r["pp"], 1) if r["pp"] else "-", round(r["tg"], 1) if r["tg"] else "-",
                     round(r["ceiling"], 1) if r["ceiling"] else "-",
                     f"{100 * r['reached']:.0f} per cent" if r["reached"] else "-"]
            print("    " + json.dumps(cells).replace('"', "'") + ",")
        print("  ]}")
        print(f"  context={{{{ hardware: '{args.hardware}', os: '{args.os}', engine: 'llama.cpp',")
        print(f"             version: '{', '.join(builds)}', model: 'as listed per row',")
        print(f"             quant: 'as listed per row', contextLength: {context},")
        print(f"             date: '{dates[-1] if dates else 'FILL-IN'}' }}}}")
        print('  status="measured"')
        print(f"  note=\"Ceiling from {bandwidth} GB/s ({source}) and each file's size times its active "
              "fraction.\" />" if bandwidth else '  note="No bandwidth given: ceilings not computed." />')
    else:
        print(f"bandwidth: {bandwidth} GB/s from {source}" if bandwidth
              else "bandwidth: none given, so no ceilings (pass --bandwidth-gbps or run predict-decode.py)")
        print("ceiling = bandwidth / (file bytes x active/total parameters); 1 GB = 1e9 bytes\n")
        header = (f"{'model file':<40} {'backend':<9} {'GB':>6} {'act GB':>6} {'pp t/s':>8} {'tg t/s':>8}"
                  f" {'tg cv%':>6} {'pp/tg':>6} {'ceiling':>8} {'reached':>7} {'gap ms':>8} {'pp TFLOP/s':>10}")
        print(header)
        print("-" * len(header))
        for r in rows:
            print(f"{r['file'][:40]:<40} {str(r['backend'])[:9]:<9} {fmt(r['file_gb'], 2, 6)} {fmt(r['active_gb'], 2, 6)}"
                  f" {fmt(r['pp'], 1)} {fmt(r['tg'], 2)} {fmt(r['tg_cv_pct'], 1, 6)} {fmt(r['pp_over_tg'], 1, 6)}"
                  f" {fmt(r['ceiling'], 1)} {fmt(r['reached'], 2, 7)} {fmt(r['unexplained_ms'], 2)}"
                  f" {fmt(r['prefill_tflops'], 2, 10)}")
            if r["id"] is None and catalogue:
                print(f"    {r['file']}: no models.json entry matches this file name; no ceiling")
            if r["reached"] and r["reached"] > 1.0:
                print("    above the ceiling: an input is wrong (bandwidth figure, file, or active bytes)")
            elif r["reached"] and r["reached"] < 1 / 3:
                print("    below a third of the ceiling: follow the challenge page's procedure")
            if r["tg_cv_pct"] and r["tg_cv_pct"] > 3.0:
                print("    generation spread above 3 per cent: something else was using the machine")
        if depth_rows:
            print(f"\n{'model file':<40} {'test':<14} {'tg t/s':>8} {'measured':>9} {'predicted':>9}")
            for d in depth_rows:
                note = "" if d["architecture"] in ("dense", None) else "  (formula assumes every layer caches every token)"
                print(f"{d['file'][:40]:<40} {d['test']:<14} {fmt(d['rate'], 2)} {fmt(d['measured_ratio'], 3, 9)}"
                      f" {fmt(d['predicted_ratio'], 3, 9)}{note}")

    if args.record:
        with notebook.open("a", encoding="utf-8") as handle:
            handle.write(json.dumps({"lab": "part-06/compare-with-prediction", "bandwidth_gbps": bandwidth,
                                     "bandwidth_source": source if bandwidth else None,
                                     "rows": rows, "depth_rows": depth_rows}) + "\n")
        print(f"\nrecorded in {notebook}", file=sys.stderr)
    return 0


if __name__ == "__main__":
    raise SystemExit(main())