"""Work out which course reference models fit in a memory budget at a target context length.

Purpose: apply the memory arithmetic from Part 4 to one machine: weights from bytes per
         parameter (or from the published GGUF file size where one is known), key-value
         cache from the model's own config.json shape, plus a reserve for the operating
         system and the engine. Prints a fit table and a suggested shortlist, and can
         append the result to the lab notebook as one JSON line.
Platform: all (pure Python 3.9+, standard library only; no accelerator needed)
Minimum memory: 8 GB
Assumes: nothing is downloaded and nothing is loaded. Every layer, key-value head and
         head dimension below was read from the model's config.json on the Hugging Face
         Hub on 2026-09-08, and every byte count marked "published" is the size of a real
         file in the GGUF repository named beside it, read from that repository's Hub file
         listing on 2026-09-13. The GGUF repositories are the ones the course model
         reference records: unsloth for Qwen3, ggml-org for gpt-oss, bartowski for Llama.

Usage: python3 fit-models.py --memory-gb 16 --context 8192 [--kv-bits 16]
                             [--reserve-gb 1.5] [--labbook labbook.md] [--json]
       python3 fit-models.py --memory-gb 16 --context 8192 --explain qwen3-8b
"""

from __future__ import annotations

import argparse
import json
import platform
import subprocess
import sys
from datetime import date
from pathlib import Path

# Bytes per parameter, including the per-block scales that the nominal bit width leaves
# out. Derived in the "Choosing a Model for a Memory Budget" lesson from the block
# structures documented at https://huggingface.co/docs/hub/gguf, and cross-checked
# against the published file sizes below.
BYTES_PER_PARAM = {
    "bf16": 2.0,
    "q8_0": 1.0625,
    "q6_k": 0.8203,
    "q5_k_m": 0.7135,
    "q4_k_m": 0.6125,
    "iq4_xs": 0.5313,
    "mxfp4": 0.5417,
}

# Order used when reporting, heaviest first.
QUANT_ORDER = ["bf16", "q8_0", "q6_k", "q5_k_m", "q4_k_m", "iq4_xs", "mxfp4"]

GB = 1_000_000_000

# The course reference set. "layers", "kv_heads" and "head_dim" come from each model's
# config.json; "files" maps a quantisation to the file name and published size in bytes
# in the GGUF repository named by "gguf_repo". A quantisation with no published file is
# still reported, with its size estimated from bytes per parameter.
MODELS = [
    {
        "id": "qwen3-1.7b",
        "name": "Qwen3-1.7B",
        "repo": "Qwen/Qwen3-1.7B",
        "gguf_repo": "unsloth/Qwen3-1.7B-GGUF",
        "params_total_b": 1.7,
        "params_active_b": 1.7,
        "layers": 28,
        "kv_heads": 8,
        "head_dim": 128,
        "quants": ["q8_0", "q6_k", "q5_k_m", "q4_k_m"],
        "files": {
            "q8_0": ("Qwen3-1.7B-Q8_0.gguf", 1_834_426_944),
            "q6_k": ("Qwen3-1.7B-Q6_K.gguf", 1_417_755_200),
            "q5_k_m": ("Qwen3-1.7B-Q5_K_M.gguf", 1_257_880_128),
            "q4_k_m": ("Qwen3-1.7B-Q4_K_M.gguf", 1_107_409_472),
        },
    },
    {
        "id": "qwen3-4b",
        "name": "Qwen3-4B",
        "repo": "Qwen/Qwen3-4B",
        "gguf_repo": "unsloth/Qwen3-4B-GGUF",
        "params_total_b": 4.0,
        "params_active_b": 4.0,
        "layers": 36,
        "kv_heads": 8,
        "head_dim": 128,
        "quants": ["q8_0", "q6_k", "q5_k_m", "q4_k_m"],
        "files": {
            "q8_0": ("Qwen3-4B-Q8_0.gguf", 4_280_405_792),
            "q6_k": ("Qwen3-4B-Q6_K.gguf", 3_306_261_792),
            "q5_k_m": ("Qwen3-4B-Q5_K_M.gguf", 2_889_514_272),
            "q4_k_m": ("Qwen3-4B-Q4_K_M.gguf", 2_497_281_312),
        },
    },
    {
        "id": "qwen3-8b",
        "name": "Qwen3-8B",
        "repo": "Qwen/Qwen3-8B",
        "gguf_repo": "unsloth/Qwen3-8B-GGUF",
        "params_total_b": 8.2,
        "params_active_b": 8.2,
        "layers": 36,
        "kv_heads": 8,
        "head_dim": 128,
        "quants": ["q8_0", "q6_k", "q5_k_m", "q4_k_m"],
        "files": {
            "q8_0": ("Qwen3-8B-Q8_0.gguf", 8_709_519_168),
            "q6_k": ("Qwen3-8B-Q6_K.gguf", 6_725_900_096),
            "q5_k_m": ("Qwen3-8B-Q5_K_M.gguf", 5_851_113_280),
            "q4_k_m": ("Qwen3-8B-Q4_K_M.gguf", 5_027_784_512),
        },
    },
    {
        # The course's cross-engine comparison model. Meta publishes no GGUF, so the
        # files are bartowski's community conversion (card: llama.cpp release b3472,
        # importance-matrix quantisation). Layer, head and head-dimension values are
        # from the course model reference; the original repository is gated.
        "id": "llama-3.1-8b",
        "name": "Llama-3.1-8B",
        "repo": "meta-llama/Llama-3.1-8B-Instruct",
        "gguf_repo": "bartowski/Meta-Llama-3.1-8B-Instruct-GGUF",
        "params_total_b": 8.0,
        "params_active_b": 8.0,
        "layers": 32,
        "kv_heads": 8,
        "head_dim": 128,
        "quants": ["q8_0", "q6_k", "q5_k_m", "q4_k_m"],
        "files": {
            "q8_0": ("Meta-Llama-3.1-8B-Instruct-Q8_0.gguf", 8_540_775_840),
            "q6_k": ("Meta-Llama-3.1-8B-Instruct-Q6_K.gguf", 6_596_011_424),
            "q5_k_m": ("Meta-Llama-3.1-8B-Instruct-Q5_K_M.gguf", 5_732_992_416),
            "q4_k_m": ("Meta-Llama-3.1-8B-Instruct-Q4_K_M.gguf", 4_920_739_232),
        },
        "note": "community conversion (bartowski); Llama 3.1 Community License; optional",
    },
    {
        "id": "qwen3-14b",
        "name": "Qwen3-14B",
        "repo": "Qwen/Qwen3-14B",
        "gguf_repo": "unsloth/Qwen3-14B-GGUF",
        "params_total_b": 14.8,
        "params_active_b": 14.8,
        "layers": 40,
        "kv_heads": 8,
        "head_dim": 128,
        "quants": ["q8_0", "q6_k", "q5_k_m", "q4_k_m"],
        "files": {
            "q8_0": ("Qwen3-14B-Q8_0.gguf", 15_698_534_784),
            "q6_k": ("Qwen3-14B-Q6_K.gguf", 12_121_938_304),
            "q5_k_m": ("Qwen3-14B-Q5_K_M.gguf", 10_514_570_624),
            "q4_k_m": ("Qwen3-14B-Q4_K_M.gguf", 9_001_753_984),
        },
    },
    {
        "id": "qwen3-30b-a3b",
        "name": "Qwen3-30B-A3B",
        "repo": "Qwen/Qwen3-30B-A3B",
        "gguf_repo": "unsloth/Qwen3-30B-A3B-GGUF",
        "params_total_b": 30.5,
        "params_active_b": 3.3,
        "layers": 48,
        "kv_heads": 4,
        "head_dim": 128,
        "quants": ["q8_0", "q6_k", "q5_k_m", "q4_k_m"],
        "files": {
            "q8_0": ("Qwen3-30B-A3B-Q8_0.gguf", 32_483_932_736),
            "q6_k": ("Qwen3-30B-A3B-Q6_K.gguf", 25_092_532_800),
            "q5_k_m": ("Qwen3-30B-A3B-Q5_K_M.gguf", 21_725_581_888),
            "q4_k_m": ("Qwen3-30B-A3B-Q4_K_M.gguf", 18_556_686_912),
        },
    },
    {
        "id": "qwen3-32b",
        "name": "Qwen3-32B",
        "repo": "Qwen/Qwen3-32B",
        "gguf_repo": "unsloth/Qwen3-32B-GGUF",
        "params_total_b": 32.8,
        "params_active_b": 32.8,
        "layers": 64,
        "kv_heads": 8,
        "head_dim": 128,
        "quants": ["q8_0", "q6_k", "q5_k_m", "q4_k_m"],
        "files": {
            "q8_0": ("Qwen3-32B-Q8_0.gguf", 34_817_719_968),
            "q6_k": ("Qwen3-32B-Q6_K.gguf", 26_883_307_168),
            "q5_k_m": ("Qwen3-32B-Q5_K_M.gguf", 23_214_832_288),
            "q4_k_m": ("Qwen3-32B-Q4_K_M.gguf", 19_762_150_048),
        },
    },
    {
        "id": "gpt-oss-20b",
        "name": "gpt-oss-20b",
        "repo": "openai/gpt-oss-20b",
        "gguf_repo": "ggml-org/gpt-oss-20b-GGUF",
        "params_total_b": 21.0,
        "params_active_b": 3.6,
        "layers": 24,
        "kv_heads": 8,
        "head_dim": 64,
        "kv_is_upper_bound": True,
        "quants": ["mxfp4"],
        "files": {"mxfp4": ("gpt-oss-20b-MXFP4.gguf", 12_109_566_624)},
    },
    {
        "id": "gpt-oss-120b",
        "name": "gpt-oss-120b",
        "repo": "openai/gpt-oss-120b",
        "gguf_repo": "ggml-org/gpt-oss-120b-GGUF",
        "params_total_b": 117.0,
        "params_active_b": 5.1,
        "layers": 36,
        "kv_heads": 8,
        "head_dim": 64,
        "kv_is_upper_bound": True,
        "quants": ["mxfp4"],
        "files": {"mxfp4": ("gpt-oss-120b-MXFP4.gguf", 63_387_346_208)},
    },
]


def kv_bytes_per_token(model: dict, kv_bits: int) -> int:
    """2 x layers x kv_heads x head_dim x bytes per element - the formula from the lesson."""
    bytes_per_element = kv_bits / 8
    return int(2 * model["layers"] * model["kv_heads"] * model["head_dim"] * bytes_per_element)


def weight_bytes(model: dict, quant: str) -> tuple[int, str]:
    """Published file size where the course has read one, else parameters x bytes per parameter."""
    published = model.get("files", {}).get(quant)
    if published is not None:
        return published[1], "published"
    per_param = BYTES_PER_PARAM[quant]
    return int(model["params_total_b"] * 1e9 * per_param), "estimated"


def default_reserve_gb(memory_gb: float) -> float:
    """The headroom rule from the lesson: OS, engine buffers and other applications."""
    if memory_gb <= 16:
        return 1.5
    if memory_gb <= 32:
        return 3.0
    if memory_gb <= 64:
        return 8.0
    return 10.0


def detect_memory_gb() -> float | None:
    """Best-effort total memory, so the script can suggest a value. Always overridable."""
    system = platform.system()
    if system == "Linux":
        try:
            for line in Path("/proc/meminfo").read_text(encoding="utf-8").splitlines():
                if line.startswith("MemTotal:"):
                    return int(line.split()[1]) * 1024 / GB
        except OSError:
            return None
    if system == "Darwin":
        try:
            out = subprocess.run(
                ["sysctl", "-n", "hw.memsize"], capture_output=True, text=True, check=True
            )
            return int(out.stdout.strip()) / GB
        except (OSError, subprocess.CalledProcessError, ValueError):
            return None
    return None


def evaluate(memory_gb: float, context: int, kv_bits: int, reserve_gb: float) -> list[dict]:
    """One row per model and quantisation, with weights, cache and whether it fits."""
    budget = memory_gb - reserve_gb
    rows = []
    for model in MODELS:
        per_token = kv_bytes_per_token(model, kv_bits)
        kv_gb = per_token * context / GB
        for quant in model["quants"]:
            raw, source = weight_bytes(model, quant)
            weights_gb = raw / GB
            total_gb = weights_gb + kv_gb
            file_name = model.get("files", {}).get(quant, (None, None))[0]
            rows.append(
                {
                    "model": model["id"],
                    "name": model["name"],
                    "quant": quant,
                    "weights_gb": round(weights_gb, 2),
                    "weights_source": source,
                    "kv_bytes_per_token": per_token,
                    "kv_gb": round(kv_gb, 2),
                    "kv_is_upper_bound": bool(model.get("kv_is_upper_bound")),
                    "total_gb": round(total_gb, 2),
                    "fits": total_gb <= budget,
                    "params_total_b": model["params_total_b"],
                    "params_active_b": model["params_active_b"],
                    "gguf_repo": model["gguf_repo"],
                    "file": file_name,
                    "downloadable": file_name is not None,
                    "note": model.get("note"),
                }
            )
    return rows


def shortlist(rows: list[dict]) -> list[dict]:
    """The heaviest quantisation that fits and can actually be downloaded, per model."""
    best: dict[str, dict] = {}
    for row in rows:
        if not row["fits"] or not row["downloadable"]:
            continue
        current = best.get(row["model"])
        if current is None or row["weights_gb"] > current["weights_gb"]:
            best[row["model"]] = row
    return sorted(best.values(), key=lambda r: r["params_total_b"])


def print_table(rows: list[dict], budget_gb: float) -> None:
    header = f"{'model':16s} {'quant':7s} {'weights':>9s} {'KV':>8s} {'total':>8s}  fit"
    print(header)
    print("-" * len(header))
    for row in rows:
        mark = "*" if row["weights_source"] == "estimated" else " "
        bound = "<" if row["kv_is_upper_bound"] else " "
        verdict = "yes" if row["fits"] else "no"
        print(
            f"{row['name']:16s} {row['quant']:7s} "
            f"{row['weights_gb']:8.2f}{mark} {row['kv_gb']:7.2f}{bound} "
            f"{row['total_gb']:8.2f}  {verdict}"
        )
    print()
    print(f"budget after reserve: {budget_gb:.2f} GB")
    if any(row["weights_source"] == "estimated" for row in rows):
        print("* weights estimated from bytes per parameter; no published file size was read")
    if any(row["kv_is_upper_bound"] for row in rows):
        print("< key-value cache is an upper bound: some layers use a sliding window")


def explain(model: dict, memory_gb: float, reserve_gb: float, context: int, kv_bits: int) -> None:
    """Show the arithmetic behind one model's rows, so the table can be checked by hand."""
    per_token = kv_bytes_per_token(model, kv_bits)
    budget = memory_gb - reserve_gb
    print(f"{model['name']}  ({model['repo']}; GGUF files from {model['gguf_repo']})")
    print(f"  config.json: num_hidden_layers {model['layers']}, num_key_value_heads "
          f"{model['kv_heads']}, head_dim {model['head_dim']}")
    print(f"  key-value cache per token = 2 x {model['layers']} x {model['kv_heads']} x "
          f"{model['head_dim']} x {kv_bits / 8:g} bytes = {per_token:,} bytes")
    for ctx in sorted({4096, 8192, 32768, context}):
        print(f"    x {ctx:>6} tokens = {per_token * ctx / GB:6.2f} GB")
    print("  weights:")
    for quant in model["quants"]:
        raw, source = weight_bytes(model, quant)
        bits = raw * 8 / (model["params_total_b"] * 1e9)
        print(f"    {quant:7s} {raw:>15,} bytes = {raw / GB:6.2f} GB  {source:9s} "
              f"({bits:.2f} bits per weight over {model['params_total_b']} B parameters)")
    print(f"  budget: {memory_gb} GB - {reserve_gb} GB reserve = {budget:.2f} GB")
    kv_gb = per_token * context / GB
    for quant in model["quants"]:
        raw, _ = weight_bytes(model, quant)
        total = raw / GB + kv_gb
        verdict = "fits" if total <= budget else "does not fit"
        print(f"    {quant:7s} at {context} tokens: {raw / GB:.2f} + {kv_gb:.2f} = "
              f"{total:.2f} GB  {verdict}")
    if model.get("kv_is_upper_bound"):
        print("  the cache figure is an upper bound: some layers use a sliding window")
    if model.get("note"):
        print(f"  note: {model['note']}")


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__.splitlines()[0])
    parser.add_argument(
        "--memory-gb",
        type=float,
        default=None,
        help="memory available to the model in GB: total unified memory, or VRAM on a discrete GPU",
    )
    parser.add_argument("--context", type=int, default=8192, help="target context length in tokens")
    parser.add_argument(
        "--kv-bits", type=int, default=16, choices=[16, 8, 4], help="key-value cache element width"
    )
    parser.add_argument(
        "--reserve-gb",
        type=float,
        default=None,
        help="memory held back for the OS, the engine and other applications",
    )
    parser.add_argument("--labbook", default=None, help="append one JSON line per run to this file")
    parser.add_argument("--json", action="store_true", help="print the full result as JSON")
    parser.add_argument(
        "--explain",
        metavar="MODEL_ID",
        default=None,
        help="print the arithmetic for one model id (for example qwen3-8b) and exit",
    )
    args = parser.parse_args()

    memory_gb = args.memory_gb
    if memory_gb is None:
        detected = detect_memory_gb()
        if detected is None:
            parser.error("could not detect memory on this platform; pass --memory-gb explicitly")
        memory_gb = round(detected, 1)
        print(f"detected total memory: {memory_gb} GB (override with --memory-gb)", file=sys.stderr)

    reserve_gb = args.reserve_gb if args.reserve_gb is not None else default_reserve_gb(memory_gb)
    budget_gb = memory_gb - reserve_gb
    if budget_gb <= 0:
        parser.error(f"reserve of {reserve_gb} GB leaves nothing of {memory_gb} GB")

    if args.explain:
        wanted = [m for m in MODELS if m["id"] == args.explain]
        if not wanted:
            parser.error(f"unknown model id {args.explain!r}; known: "
                         + ", ".join(m["id"] for m in MODELS))
        explain(wanted[0], memory_gb, reserve_gb, args.context, args.kv_bits)
        return

    rows = evaluate(memory_gb, args.context, args.kv_bits, reserve_gb)
    picks = shortlist(rows)
    download_gb = round(sum(p["weights_gb"] for p in picks), 2)

    result = {
        "lab": "part-04/fit-models",
        "date": date.today().isoformat(),
        "memory_gb": memory_gb,
        "reserve_gb": reserve_gb,
        "budget_gb": round(budget_gb, 2),
        "context": args.context,
        "kv_bits": args.kv_bits,
        "shortlist": [
            {
                "model": p["model"],
                "quant": p["quant"],
                "weights_gb": p["weights_gb"],
                "kv_gb": p["kv_gb"],
                "total_gb": p["total_gb"],
                "gguf_repo": p["gguf_repo"],
                "file": p["file"],
                "note": p.get("note"),
            }
            for p in picks
        ],
        "download_gb": download_gb,
    }

    if args.json:
        print(json.dumps(result, indent=2))
    else:
        print(f"memory {memory_gb} GB, reserve {reserve_gb} GB, context {args.context} tokens, "
              f"KV at {args.kv_bits}-bit")
        print()
        print_table(rows, budget_gb)
        print()
        if picks:
            print("shortlist (heaviest quantisation that fits, per model):")
            for p in picks:
                print(f"  {p['name']:16s} {p['quant']:7s} {p['weights_gb']:8.2f} GB   "
                      f"{p['gguf_repo']}  {p['file']}")
                if p.get("note"):
                    print(f"  {'':16s} {'':7s} {'':11s} {p['note']}")
            print(f"  total download if you took every row: {download_gb:.2f} GB")
            print("  this is a menu, not a shopping list: the lab downloads one or two of these")
        else:
            print("nothing in the reference set fits. Try a shorter context or --kv-bits 8.")

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


if __name__ == "__main__":
    main()