#!/usr/bin/env python3
"""grams-to-imatrix.py — the engine's own calibration recording, re-voiced as a
llama.cpp importance matrix for `rad-convert --imatrix F`. RUNS ON DANKSERVER
beside the recording. Python 3 stdlib only: numpy must NOT be required, because
this runs in the plain container next to a 335 GiB checkpoint, not in a venv.

WHY. `rad-convert` takes its per-column activation second moments and its
expert-popularity profile from ONE input, an imatrix (core/format/imatrix.h).
Only the GGUF form carries a per-expert dimension — with the legacy flat form,
or with none, every expert in a layer is weighted by the layer's aggregate,
which is exactly the wrong thing on a 512-expert MoE. The calibration tap has
already measured everything; this tool only re-voices it, so the imatrix and the
protected-expert census (derive-protected-experts.py) are the SAME recording:
one measurement, two consumers, neither able to disagree about which experts are
hot.

INPUT. RADGRAM2 files in a RADIANCE_CALIB_DIR output directory
(arch/common/rad_arch.h, calib_write). Header law, same parse as
scripts/derive-protected-experts.py — 32 bytes, `<8sIIQI`: magic `RADGRAM2`,
u32 k, u32 rank, u64 rows, u32 domain, a pad word, then k*k f32 row-major. Two
kinds of name, both GGUF-style, both with the checkpoint's GLOBAL expert
numbering (rad_block_moe_fp8.h: ge = expert_base + e):

    gram.r<rank>.blk.L.ffn_down_exps.bin        the layer's pooled gram
    gram.r<rank>.blk.L.ffn_down_exps.<e>.bin    expert e's own gram

`rows` is that file's routing/token count and the matrix is Σ x xᵀ over those
rows, so **the diagonal is exactly the per-column activation second moment an
imatrix wants**, and `rows` is the count the reader hands back from `count()`.
Every off-diagonal term is dropped on purpose: the GGUF format has nowhere to
put a full matrix, and the quantiser asks for columns.

OUTPUT, per layer L (core/format/imatrix.h: names WITHOUT `.weight`, a MoE
layer's experts folded into ONE entry, so expert 9 is slice 9 of
`blk.L.ffn_down_exps`):

  blk.L.ffn_down_exps    in_sum2 [k, n_expert] expert-major, counts [1, n_expert]
  blk.L.ffn_gate_up_exps the SAME matrices at the SAME offsets. A routed expert's
                         gate_up input is the same hidden-state rows as its down
                         input, so the second moment over them is identical and
                         so is the routing count. Two tensor infos over one
                         payload: the reader maps each half independently
                         (imatrix.cpp joins `<base>.in_sum2` with `<base>.counts`
                         by base name), and copying ~200 MB to state one fact
                         twice buys no correctness.
  blk.L.ssm_inz          [k] and [1], from the pooled gram: that row's input is
                         the same hidden state, so its columns take the pooled
                         diagonal and its count is the layer's pooled rows.

  Then STOP. Nothing whose input width is not the gram's k — no ssm_out, no
  attn_output, no mtp.*, no ple_ngram. importance_for() answers false when the
  entry's cols do not equal the weight's width and the convert proceeds without
  it: an entry the reader silently ignores is a lie in the artifact, worse than
  an absent one. The corollary to check before trusting the two hidden-width
  entries: the tap writes the gate_up gram at g.n_embd and the down gram at the
  expert width, so the gram's k must equal the checkpoint's hidden width (its
  config.json). Where they differ, gate_up/ssm_inz are inert, not wrong.

INHERITANCE. An expert with no own gram carries the layer's POOLED diagonal, and
its count stays its own routing rows — the rule the engine itself applies when it
picks a Hessian (libquant/lq_hess.h, calib_resolve: an expert's own gram only
once its rank files hold kOwnRowsPerK * K rows; below that a fit to so few rows
is worse than round-to-nearest, and those experts carry a fraction of a percent
of the routing). The count never inherits: popularity is `count`'s job. The count
of inherited experts is reported on stderr.

RANK FOLDING. TP=2 writes gram.r0.* and gram.r1.* with GLOBAL expert numbering,
so the ranks describe DISJOINT experts and the fold is a UNION into the one
[k, n_expert] entry. A pooled gram is the opposite case: every rank summed the
same stem over its own rows, so the fold is a SUM, rows included (calib_sum,
calib_rows). The same global expert in two per-expert files is legal only when
the numbers agree exactly; otherwise the recording is corrupt and the entry would
be whichever file sorted last, so refuse.

GGUF v3 (the only version imatrix.cpp accepts — version != 3 is a refusal):
`general.alignment` u32, `imatrix.chunk_size` u32, `imatrix.chunk_count` u32,
`imatrix.datasets` array of strings; tensors f32 (ggml type 0) with offsets
relative to the aligned data section. Every tensor offset is padded up to
`general.alignment` and the data section itself starts aligned. That is the
classic hand-rolled-writer bug: an unaligned offset makes the reader's extent
arithmetic point into the middle of the previous tensor's floats, and it reads
them without complaint.

chunk_count DERIVATION. The corpus row count is the LARGEST single-layer pooled
row total, not a sum over layers: every layer's tap accumulated over the whole
corpus, so adding the layers counts each token once per layer and inflates the
figure by the layer count. chunk_count = that row total // chunk_size, the
"chunks of chunk_size" line imatrix.cpp prints for itself.

REFUSALS, each exit 2 with a message naming what is missing:
  NO_GRAMS         the directory holds no RADGRAM2 down-projection gram
  NOT_RADGRAM2     a gram whose magic is not RADGRAM2
  SHORT            a gram whose length is not 32 + k*k*4
  NO_POOLED        a layer with expert grams and no pooled gram to inherit from
  K_MISMATCH       a layer whose files disagree about k (rank pooled grams, or an
                   expert gram against its layer's pooled gram)
  EXPERT_CONFLICT  a global expert recorded twice with different numbers
  EXISTS           the -o path already holds a file (the trash law: move it)
  CHUNK_COUNT      the derived chunk count does not fit the u32 metadata field

USAGE. grams-to-imatrix.py <calib-dir> -o <out.gguf>
       [--datasets NAME,...] [--chunk-size N] [--min-rows N] [-v]
"""
import argparse
import os
import re
import struct
import sys

MAGIC = b"RADGRAM2"
HDR = 32                       # 8s I I Q I + pad, rad_arch.h
KOWN_ROWS_PER_K = 4            # libquant/lq_hess.h kOwnRowsPerK
ALIGNMENT = 32                 # general.alignment, and the pad unit

POOLED_RE = re.compile(r"^gram\.r(\d+)\.blk\.(\d+)\.ffn_down_exps\.bin$")
EXPERT_RE = re.compile(r"^gram\.r(\d+)\.blk\.(\d+)\.ffn_down_exps\.(\d+)\.bin$")

GT_U32, GT_STR, GT_ARR = 4, 8, 9
GGML_F32 = 0


def refuse(msg):
    sys.stderr.write("REFUSED %s\n" % msg)
    sys.exit(2)


def read_gram(path):
    """(k, rank, rows, domain, diagonal) from one RADGRAM2 file."""
    try:
        with open(path, "rb") as f:
            raw = f.read()
    except OSError as e:
        refuse("UNREADABLE %s: %s" % (path, e))
    if len(raw) < HDR:
        refuse("SHORT %s: %d bytes, under the %d-byte RADGRAM2 header"
               % (path, len(raw), HDR))
    magic, k, rank, rows, domain = struct.unpack_from("<8sIIQI", raw, 0)
    if magic != MAGIC:
        refuse("NOT_RADGRAM2 %s: magic %r, expected b'RADGRAM2'" % (path, magic))
    if k <= 0:
        refuse("SHORT %s: header claims k=%d, which describes no matrix" % (path, k))
    need = HDR + k * k * 4
    if len(raw) != need:
        refuse("SHORT %s: %d bytes, expected %d (32 + k*k*4 for k=%d)"
               % (path, len(raw), need, k))
    # The diagonal only, read by stride out of the row-major matrix.
    diag = [struct.unpack_from("<f", raw, HDR + (i * k + i) * 4)[0]
            for i in range(k)]
    return k, rank, rows, domain, diag


def load_dir(d, verbose):
    """{layer: {"pooled": gram|None, "experts": {global e: gram}}} + rotated count.

    Files outside the ffn_down_exps stem are not ours to fold: the tap names a
    pooled gram after every projection it taps, and those carry other widths."""
    if not os.path.isdir(d):
        refuse("NO_GRAMS %s: not a directory" % d)
    layers, rotated = {}, 0
    for name in sorted(os.listdir(d)):
        m = POOLED_RE.match(name) or EXPERT_RE.match(name)
        if not m:
            continue
        path = os.path.join(d, name)
        if not os.path.isfile(path):
            continue
        pooled = bool(POOLED_RE.match(name))
        layer = int(m.group(2))
        k, rank, rows, domain, diag = read_gram(path)
        if domain != 0:
            rotated += 1
        L = layers.setdefault(layer, {"pooled": None, "experts": {}})
        if pooled:
            P = L["pooled"]
            if P is None:
                L["pooled"] = {"k": k, "rows": rows, "diag": diag, "paths": [path]}
            elif P["k"] != k:
                refuse("K_MISMATCH layer %d: pooled grams %s (k=%d) and %s (k=%d)"
                       " disagree about the width"
                       % (layer, P["paths"][0], P["k"], path, k))
            else:
                # calib_sum: each rank summed the same stem over its own rows.
                P["rows"] += rows
                P["paths"].append(path)
                for i in range(k):
                    P["diag"][i] += diag[i]
            if verbose:
                sys.stderr.write("pooled layer %d: %s k=%d rows=%d\n"
                                 % (layer, name, k, rows))
            continue
        e = int(m.group(3))
        own = L["experts"].get(e)
        if own is None:
            L["experts"][e] = {"k": k, "rows": rows, "diag": diag,
                               "path": path, "ranks": [rank]}
            if verbose:
                sys.stderr.write("expert layer %d r%d e%d: rows=%d\n"
                                 % (layer, rank, e, rows))
            continue
        if own["rows"] != rows or own["diag"] != diag:
            refuse("EXPERT_CONFLICT layer %d expert %d: %s (rows=%d) and %s"
                   " (rows=%d) record the same global expert with different"
                   " numbers" % (layer, e, own["path"], own["rows"], path, rows))
        own["ranks"].append(rank)
        sys.stderr.write("note: layer %d expert %d recorded identically by rank %s"
                         " and %s; folded once\n"
                         % (layer, e, own["path"], name))
    if not layers:
        refuse("NO_GRAMS %s: no RADGRAM2 gram.r<rank>.blk.<L>.ffn_down_exps"
               "[.<expert>].bin files in the directory" % d)
    return layers, rotated


def build_entries(layers, min_rows, verbose):
    """([ (name, dims, payload|None, key) ], census, layers emitted)."""
    entries = []
    census = {"inherited": 0, "experts": 0, "routing": 0, "rows": 0, "layers": 0}
    for layer in sorted(layers):
        P, E = layers[layer]["pooled"], layers[layer]["experts"]
        if E and P is None:
            refuse("NO_POOLED layer %d: %d expert gram(s) but no pooled gram (an"
                   " expected gram.r<rank>.blk.%d.ffn_down_exps.bin), so an expert"
                   " under the row threshold would have nothing to inherit"
                   % (layer, len(E), layer))
        if P is None:
            continue
        k = P["k"]
        census["layers"] += 1
        census["rows"] = max(census["rows"], P["rows"])
        dn = "blk.%d.ffn_down_exps" % layer
        # A staging key names ONE payload. A gate_up entry reuses the down
        # layer's bytes by carrying the same key with no payload of its own;
        # in_sum2 and counts are different bytes and must never share a key.
        ksum, kcnt = (layer, "sum2"), (layer, "counts")
        kpgm, kpcnt = (layer, "pooled_sum2"), (layer, "pooled_counts")
        # One entry holds one width: an expert gram recorded at another k would
        # inherit (or be inherited by) a diagonal of the wrong length, and the
        # reader would then answer `cols` from whichever file set dims[0].
        for e in sorted(E):
            if E[e]["k"] != k:
                refuse("K_MISMATCH layer %d expert %d: %s holds k=%d and this"
                       " layer's pooled gram is k=%d" % (layer, e, E[e]["path"],
                                                          E[e]["k"], k))

        if not E:
            # No per-expert recording for this layer: emit the dense form
            # ([k], [1]) rather than 512 identical slices, which would claim a
            # per-expert measurement that was never taken.
            entries.append((dn + ".in_sum2", [k],
                            struct.pack("<%df" % k, *P["diag"]), ksum))
            entries.append((dn + ".counts", [1],
                            struct.pack("<f", float(P["rows"])), kcnt))
            sys.stderr.write("note: layer %d has no per-expert grams; %s is"
                             " dense over the pooled gram\n" % (layer, dn))
        else:
            floor = min_rows if min_rows > 0 else KOWN_ROWS_PER_K * k
            n_expert = max(E) + 1
            sum2, counts = bytearray(), bytearray()
            for e in range(n_expert):
                own = E.get(e)
                if own is not None and own["rows"] >= floor:
                    src = own["diag"]
                    rows = own["rows"]
                else:
                    src = P["diag"]
                    rows = own["rows"] if own is not None else 0
                    census["inherited"] += 1
                sum2 += struct.pack("<%df" % k, *src)
                counts += struct.pack("<f", float(rows))
                census["experts"] += 1
                census["routing"] += rows
            if min(E) != 0:
                sys.stderr.write("note: layer %d has no gram for experts 0..%d;"
                                 " those slices carry the pooled gram at count 0\n"
                                 % (layer, min(E) - 1))
            entries.append((dn + ".in_sum2", [k, n_expert], bytes(sum2), ksum))
            entries.append((dn + ".counts", [1, n_expert], bytes(counts), kcnt))
            gu = "blk.%d.ffn_gate_up_exps" % layer
            entries.append((gu + ".in_sum2", [k, n_expert], None, ksum))
            entries.append((gu + ".counts", [1, n_expert], None, kcnt))
            if verbose:
                sys.stderr.write("layer %d: k=%d n_expert=%d pooled_rows=%d"
                                 " own-gram floor=%d\n"
                                 % (layer, k, n_expert, P["rows"], floor))

        inz = "blk.%d.ssm_inz" % layer
        entries.append((inz + ".in_sum2", [k],
                        struct.pack("<%df" % k, *P["diag"]), kpgm))
        entries.append((inz + ".counts", [1],
                        struct.pack("<f", float(P["rows"])), kpcnt))
    return entries, census


def enc_str(v):
    b = v.encode("utf-8")
    return struct.pack("<Q", len(b)) + b


def write_gguf(path, entries, datasets, chunk_size, chunk_count):
    """Write GGUF v3 and return its size in bytes.

    One payload per key: entries that share a key share the bytes and the
    offset, so a second tensor info states the same fact without copying it."""
    kv = bytearray()
    kv += enc_str("general.alignment") + struct.pack("<I", GT_U32) \
        + struct.pack("<I", ALIGNMENT)
    kv += enc_str("imatrix.chunk_size") + struct.pack("<I", GT_U32) \
        + struct.pack("<I", chunk_size)
    kv += enc_str("imatrix.chunk_count") + struct.pack("<I", GT_U32) \
        + struct.pack("<I", chunk_count)
    kv += enc_str("imatrix.datasets") + struct.pack("<I", GT_ARR) \
        + struct.pack("<I", GT_STR) + struct.pack("<Q", len(datasets))
    for d in datasets:
        kv += enc_str(d)

    head = bytearray()
    head += b"GGUF" + struct.pack("<I", 3)
    head += struct.pack("<Q", len(entries)) + struct.pack("<Q", 4)
    head += kv

    payloads = {}
    for name, _dims, payload, key in entries:
        if payload is None:
            continue
        prev = payloads.get(key)
        if prev is not None and prev != payload:
            # Two different payloads on one key would ship whichever was
            # appended last while BOTH tensor infos pointed at it: the counts
            # payload landing over the in_sum2 payload is exactly how a
            # hand-rolled GGUF writer loses half its data without an error.
            refuse("INTERNAL %s: two different payloads share one staging key"
                   % name)
        payloads.setdefault(key, payload)
    infos = bytearray()
    offsets, order, size = {}, [], 0
    for name, dims, payload, key in entries:
        if key not in payloads:
            refuse("INTERNAL %s: no payload staged for its key" % name)
        if key not in offsets:
            offsets[key] = (size + ALIGNMENT - 1) // ALIGNMENT * ALIGNMENT
            size = offsets[key] + len(payloads[key])
            order.append(key)
        infos += enc_str(name) + struct.pack("<I", len(dims))
        for d in dims:
            infos += struct.pack("<Q", d)
        infos += struct.pack("<I", GGML_F32) + struct.pack("<Q", offsets[key])

    body = bytes(head) + bytes(infos)
    data_start = (len(body) + ALIGNMENT - 1) // ALIGNMENT * ALIGNMENT
    with open(path, "wb") as f:
        f.write(body)
        f.write(b"\0" * (data_start - len(body)))
        # The offsets in the tensor infos are the file's law, not a hint: each
        # payload must land at its own padded offset. Writing them back to back
        # instead leaves every reader that trusts the info — imatrix.cpp does —
        # reading its neighbour's floats: plausible values, no error, wrong
        # quantisation.
        here = data_start
        for key in order:
            f.write(b"\0" * (data_start + offsets[key] - here))
            f.write(payloads[key])
            here = data_start + offsets[key] + len(payloads[key])
        return f.tell()


def main():
    ap = argparse.ArgumentParser(
        description="Fold a radiance calibration recording into a GGUF v3 imatrix.")
    ap.add_argument("calib_dir")
    ap.add_argument("-o", "--out", required=True)
    ap.add_argument("--datasets", default="",
                    help="comma-separated names for imatrix.datasets")
    ap.add_argument("--chunk-size", type=int, default=4096,
                    help="rows per chunk for imatrix.chunk_size (default 4096)")
    ap.add_argument("--min-rows", type=int, default=0,
                    help="own-gram row threshold (default 4*k, the engine's"
                         " kOwnRowsPerK law)")
    ap.add_argument("-v", "--verbose", action="store_true")
    a = ap.parse_args()

    if os.path.exists(a.out):
        # The repo's deletion law: an artifact that exists is moved to trash by
        # a human, never silently overwritten by a re-run of the tool that made
        # it, so a second convert cannot destroy the imatrix the first one fed.
        refuse("EXISTS %s: an imatrix is already there; move it to"
               " .AGENTTRASH/ or name another -o" % a.out)
    if a.chunk_size <= 0:
        refuse("BAD_CHUNK_SIZE --chunk-size %d is not positive" % a.chunk_size)
    layers, rotated = load_dir(a.calib_dir, a.verbose)
    entries, census = build_entries(layers, a.min_rows, a.verbose)
    if not entries:
        refuse("NO_GRAMS %s: no layer holds a pooled gram to fold" % a.calib_dir)
    if rotated:
        sys.stderr.write("note: %d gram(s) recorded in a rotated domain: their"
                         " diagonal is the rotated activation's second moment,"
                         " not the model's own\n" % rotated)

    chunk_count = census["rows"] // a.chunk_size
    if chunk_count > 0xFFFFFFFF:
        # imatrix.chunk_count is a u32 in the reader's grammar, so there is no
        # legal way to record a bigger one. Reaching this means either a chunk
        # size picked for a far larger run or a row total read out of the wrong
        # header offset — rows is a u64 and a u32 parse is off by 2**32 — and
        # the reader would trust whatever number landed in the field.
        refuse("CHUNK_COUNT %d = %d rows // %d does not fit the u32"
               " imatrix.chunk_count field; raise --chunk-size"
               % (chunk_count, census["rows"], a.chunk_size))
    nbytes = write_gguf(a.out, entries, [d for d in a.datasets.split(",") if d],
                        a.chunk_size, chunk_count)
    sys.stderr.write("IMATRIX_OK entries=%d layers=%d experts=%d inherited=%d"
                     " rows=%d routing=%d chunks=%d/%d bytes=%d out=%s\n"
                     % (len(entries), census["layers"], census["experts"],
                        census["inherited"], census["rows"], census["routing"],
                        chunk_count, a.chunk_size, nbytes, a.out))


if __name__ == "__main__":
    main()