#!/usr/bin/env python3
"""
build_recipe.py — size-solve a thireus-style per-tensor quant recipe for
Qwen3.8-27B against a hard VRAM weight budget.

WHY THIS EXISTS
---------------
The thireus q6_0-7bpw map for Qwen3.6-27B is a flat list of anchored regexes
with hardcoded layer indices. Qwen3.8-27B has the SAME tensor shape (WATCH
gate G3: 866 tensors, blk.0-63 + blk.64 MTP), so the map transfers verbatim
— but "verbatim" is exactly the assumption that fails silently: a
--tensor-type pattern that matches nothing is NOT an error. So this script
does two jobs:

  1. GATE: verify every regex in the inherited map matches exactly one real
     tensor in the converted GGUF, and that every weight tensor is covered.
     Refuses to emit a recipe if not.
  2. SOLVE: the inherited map lands ~22,887 MiB (the ThinkingCap reference).
     We measured a larger budget than that on this rig, so spend the
     headroom by promoting Q6_K -> Q8_0 in a principled order until the
     budget is hit. This is the "Goldilocks" step.

PROMOTION ORDER (and why)
-------------------------
We do NOT have a fresh per-tensor sensitivity ranking for 3.8 (that needs
benchmark_each_tensor.sh, hours of GPU). So the order below is derived from
the structure of the inherited map itself, which already encodes thireus'
sensitivity findings for this architecture:

  Groups the map protects COMPLETELY at Q8_0 (attn_k/v/output, ssm_alpha,
  ssm_beta, nextn.eh_proj) are already maxed - nothing to promote.

  Groups the map protects PARTIALLY are, by construction, the ones sitting
  right at the sensitivity boundary - thireus spent bits on some layers of
  these and not others. Completing them is the highest-confidence use of
  headroom, so they go first, most-protected group first:
    ssm_out      (43/48 already q8_0 -> finish the last 5)
    attn_gate    (30/48 already q8_0 -> finish 18)
    attn_q       ( 2/17 already q8_0 -> finish 15)
    attn_qkv     ( 4/48 already q8_0 -> finish 44)
  Then the FFN bulk, ffn_down first (the map itself favours it: 9 layers at
  q8_0 vs 6 for ffn_gate and 5 for ffn_up - same ordering llama.cpp's own
  quant heuristics use):
    ffn_down -> ffn_gate -> ffn_up

Within a group we promote LOW layer index first: early layers are the
conventional sensitivity hot spot and the map's own q8_0 picks cluster there.

OUTPUT
------
A --tensor-type argument file for llama-quantize, plus a printed size
report. Nothing is quantized here.
"""

import argparse
import re
import sys
from pathlib import Path

sys.path.insert(0, r"C:\path\to\llama.cpp\gguf-py")  # point at your llama.cpp checkout
from gguf import GGUFReader  # noqa: E402

# bytes per element, by ggml type name
BPE = {
    "F32": 4.0,
    "F16": 2.0,
    "BF16": 2.0,
    "Q8_0": 34.0 / 32.0,      # 32 vals -> 2B scale + 32B data
    "Q6_K": 210.0 / 256.0,
    "Q5_K": 176.0 / 256.0,
    "Q4_K": 144.0 / 256.0,
}

# group -> promotion priority (lower = promoted sooner). See docstring.
PROMOTION_ORDER = [
    "ssm_out",
    "attn_gate",
    "attn_q",
    "attn_qkv",
    "ffn_down",
    "ffn_gate",
    "ffn_up",
]

MIB = 1024.0 * 1024.0


def stem_of(name: str) -> str:
    """blk.12.ffn_down.weight -> ffn_down"""
    m = re.match(r"^blk\.(\d+)\.(.+)\.weight$", name)
    return m.group(2) if m else name


def layer_of(name: str) -> int:
    m = re.match(r"^blk\.(\d+)\.", name)
    return int(m.group(1)) if m else -1


def main() -> int:
    ap = argparse.ArgumentParser()
    ap.add_argument("--gguf", required=True, help="converted BF16 GGUF")
    ap.add_argument("--map", required=True, help="inherited tensortypes.txt")
    ap.add_argument("--budget-mib", type=float, required=True,
                    help="hard ceiling for the quantized file, MiB")
    ap.add_argument("--out", required=True, help="tensor-type args file to write")
    args = ap.parse_args()

    print(f"reading {args.gguf} ...", flush=True)
    r = GGUFReader(args.gguf)
    tensors = {t.name: t for t in r.tensors}
    print(f"  {len(tensors)} tensors")

    # ---- load inherited map ----
    assign: dict[str, str] = {}
    unmatched: list[str] = []
    for line in Path(args.map).read_text(encoding="utf-8").splitlines():
        line = line.strip()
        if not line or line.startswith("#"):
            continue
        pat, _, ty = line.partition("=")
        # ^blk.0.ffn_down.weight$  -> literal name
        name = pat.strip().lstrip("^").rstrip("$").replace("\\.", ".")
        if name in tensors:
            assign[name] = ty.strip().upper()
        else:
            unmatched.append(name)

    # ---- GATE 1: every map entry must hit a real tensor ----
    if unmatched:
        print(f"\nFATAL: {len(unmatched)} map entries match NO tensor in this GGUF.")
        print("This is the silent-failure mode - refusing to emit a recipe.")
        for n in unmatched[:15]:
            print("   ", n)
        return 2

    # ---- GATE 2: report coverage of weight tensors ----
    quantizable = {
        n: t for n, t in tensors.items()
        if len(t.shape) >= 2 and n.endswith(".weight")
    }
    uncovered = sorted(set(quantizable) - set(assign))
    print(f"  map covers {len(assign)}/{len(quantizable)} 2-D weight tensors")
    if uncovered:
        print(f"  {len(uncovered)} NOT in map (will take llama-quantize's base type):")
        for n in uncovered[:10]:
            print("     ", n)

    def elems(name: str) -> int:
        n = 1
        for d in tensors[name].shape:
            n *= int(d)
        return n

    # ---- pin the big tensors the inherited map does NOT cover ----
    #
    # This is where a 310 MiB size-model error came from on the first pass, and
    # it matters more than it sounds: token_embd and output are 1,271,398,400
    # elements EACH here (vocab 248,320 x 5,120) - together ~2.3 GiB, roughly
    # 10% of the whole budget, and NEITHER is in thireus' map.
    #
    #   output.weight   -> Q8_0. It is the LM head; llama-quantize already
    #                      promotes it by default, and reproducing the
    #                      ThinkingCap reference size requires it (assuming
    #                      Q6_K under-predicted that file by ~293 MiB).
    #   token_embd      -> Q6_K. Embeddings tolerate it; this is what the
    #                      reference quant used. Dropping it to Q5_K would buy
    #                      ~200 MiB for FFN upgrades but deviates from the
    #                      proven recipe - not worth the risk unprompted.
    #   *ssm_conv1d     -> F32. Tiny (7.5 MiB total at F32) and llama.cpp
    #                      generally refuses to quantize conv kernels anyway.
    #                      Pinned so the output size is deterministic rather
    #                      than dependent on llama-quantize's internal rules.
    for n, t in tensors.items():
        if n in assign:
            continue
        if n == "output.weight":
            assign[n] = "Q8_0"
        elif n == "token_embd.weight":
            assign[n] = "Q6_K"
        elif "conv1d" in n:
            assign[n] = "F32"

    def total_mib(a: dict[str, str]) -> float:
        tot = 0.0
        for n, t in tensors.items():
            e = elems(n)
            if n in a:
                tot += e * BPE[a[n]]
            else:
                # remaining uncovered: 1-D/norms stay F32, others base Q6_K
                ty = "F32" if len(t.shape) < 2 else "Q6_K"
                tot += e * BPE[ty]
        return tot / MIB

    base = total_mib(assign)
    print(f"\ninherited map size:  {base:10.1f} MiB")
    print(f"budget:              {args.budget_mib:10.1f} MiB")
    print(f"headroom:            {args.budget_mib - base:10.1f} MiB")

    if base > args.budget_mib:
        print("\nWARNING: inherited map ALREADY exceeds budget - would need "
              "demotion, not promotion. Not implemented; stopping.")
        return 3

    # ---- SOLVE: promote Q6_K -> Q8_0 in principled order ----
    # token_embd is excluded from promotion on purpose: at 1.27B elements it
    # would eat ~293 MiB of headroom for the least sensitivity-per-byte of any
    # tensor in the model. That headroom belongs to the FFN stack.
    candidates = [n for n, ty in assign.items()
                  if ty == "Q6_K" and n != "token_embd.weight"]

    def rank(n: str):
        s = stem_of(n)
        gi = PROMOTION_ORDER.index(s) if s in PROMOTION_ORDER else len(PROMOTION_ORDER)
        return (gi, layer_of(n))

    candidates.sort(key=rank)

    promoted: list[str] = []
    cur = base
    for n in candidates:
        delta = elems(n) * (BPE["Q8_0"] - BPE["Q6_K"]) / MIB
        if cur + delta <= args.budget_mib:
            assign[n] = "Q8_0"
            promoted.append(n)
            cur += delta

    print(f"\npromoted {len(promoted)}/{len(candidates)} Q6_K tensors to Q8_0")
    by_stem: dict[str, int] = {}
    for n in promoted:
        by_stem[stem_of(n)] = by_stem.get(stem_of(n), 0) + 1
    for s in PROMOTION_ORDER:
        tot = sum(1 for n in candidates if stem_of(n) == s)
        if tot:
            print(f"   {s:12s} {by_stem.get(s,0):3d}/{tot:3d}")

    final = total_mib(assign)
    n_q8 = sum(1 for v in assign.values() if v == "Q8_0")
    n_q6 = sum(1 for v in assign.values() if v == "Q6_K")
    tot_elems = sum(elems(n) for n in assign)
    bpw = final * MIB * 8.0 / tot_elems

    print(f"\nfinal size:          {final:10.1f} MiB   ({final*MIB/1e9:.2f} GB decimal)")
    print(f"under budget by:     {args.budget_mib - final:10.1f} MiB")
    print(f"mapped tensors:      {n_q8} Q8_0 / {n_q6} Q6_K")
    print(f"effective bpw (mapped tensors): {bpw:.3f}")

    with open(args.out, "w", encoding="utf-8") as f:
        for n in sorted(assign):
            # anchored, dots escaped - same format as the proven thireus file.
            # anchors are NOT optional: an unanchored "attn_q" would also match
            # attn_qkv and silently mis-assign 48 tensors.
            f.write(f"^{re.escape(n)}$={assign[n].lower()}\n")
    print(f"\nwrote {args.out}")
    return 0


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