"""Derive GLM-5.3-Flash-NVFP4-FP8ATTN (r2: MLP + ATTENTION FP8) from
GLM-5.3-Flash-NVFP4 (rev caca4e6a; expert weight shards are byte-identical to
357b45cc — only the input scales and card differ between those revisions).

Targets the patched serving image that passes the real quant_config into the
KDA/MLA attention constructors (kda.py strip removed, model.py:329
None -> real). See the model card and serving/README.md.

Run convert_lmhead.py afterward to reproduce the published checkpoint
(block-FP8 lm_head) bit-for-bit.

Converted surfaces and block sizes (weight F8_E4M3 + weight_scale_inv F32,
dequant multiplier amax/448):

  FP8_BLOCK128 ([128,128], the proven MLP path):
    - *.mlp.shared_experts.{gate,up,down}_proj   layers 3..45 (incl. MTP 45)
    - *.mlp.{gate,up,down}_proj                  dense layers 0..2
    - *.self_attn.o_proj                         ALL 46 layers (KDA + MLA)
    - *.self_attn.{q_a_proj,kv_a_proj_with_mqa,q_b_proj}   MLA layers
      (q_a + kv_a fuse into fused_qkv_a_proj at runtime: offsets 0/1536 and
       sizes 1536/512 all divide 128 because qk_rope_head_dim=0)

  FP8_BLOCK32 ([32,32], new dispatch in the r2 image's modelopt.py):
    - *.self_attn.{q,k,v,b,f_a,g_a}_proj         KDA layers
      (they fuse into in_proj_qkvbfg_a whose per-rank TP2 layout is
       q|k|v 4096 each | b 32 | f_a 128 | g_a 128 — block_n must divide 32;
       square because Fp8LinearMethod's activation group = block_n)

  Kept BF16 (deliberate):
    - kv_b_proj      — MLA absorption stores BF16 W_UK/W_UV copies regardless
                       (mla_attention.py:1022); net saving ~0.09 GiB/rank and
                       _get_kv_b_proj_input_dtype has an untested fp8 branch.
    - f_b_proj/g_b_proj — delta-rule gate projections, not in the tasking.
    - indexer.*, convs, norms, embed_tokens, lm_head, visual.* — as before.

Untouched shards are HARDLINKED to the source. Manifest carries both the
main-model spelling (model.language_model.layers.N.*) and the MTP draft
runtime spelling (model.layers.45.*), for MLP AND attention.
"""

import json
import os
import re
import struct
import time

import ml_dtypes
import numpy as np

SRC = None  # set from --input in __main__
DST = None  # set from --output in __main__
FIRST_K_DENSE = 3
NUM_LAYERS = 46  # 45 decoder + MTP layer 45
NUM_EXPERTS = 288

MLP_RE = re.compile(
    r"^model\.language_model\.layers\.(\d+)\.mlp\."
    r"(shared_experts\.)?(gate|up|down)_proj\.weight$"
)
ATTN128_RE = re.compile(
    r"^model\.language_model\.layers\.(\d+)\.self_attn\."
    r"(o_proj|q_a_proj|kv_a_proj_with_mqa|q_b_proj)\.weight$"
)
ATTN32_RE = re.compile(
    r"^model\.language_model\.layers\.(\d+)\.self_attn\."
    r"(q|k|v|b|f_a|g_a)_proj\.weight$"
)

# ignore-list patterns whose modules are now FP8-quantized
UNIGNORE = {
    # MLP family (as in convert_fp8mlp.py)
    "*.mlp.shared_experts.gate_up_proj",
    "*.mlp.shared_experts.down_proj",
    "*.mlp.shared_experts.gate_proj",
    "*.mlp.shared_experts.up_proj",
    "*.mlp.gate_up_proj",
    "*.mlp.down_proj",
    "*.mlp.gate_proj",
    "*.mlp.up_proj",
    # attention: KDA fused in_proj shards + fused spellings
    "*.self_attn.q_proj",
    "*.self_attn.k_proj",
    "*.self_attn.v_proj",
    "*.self_attn.b_proj",
    "*.self_attn.f_a_proj",
    "*.self_attn.g_a_proj",
    "*.self_attn.in_proj_qkvbfg_a",
    "*.self_attn.fused_qkvbfg_a_proj",
    # attention: o_proj (KDA + MLA) and the MLA q/lora side
    "*.self_attn.o_proj",
    "*.self_attn.q_a_proj",
    "*.self_attn.kv_a_proj_with_mqa",
    "*.self_attn.fused_qkv_a_proj",
    "*.self_attn.fused_qkv_a_proj_with_mqa",
    "*.self_attn.q_b_proj",
}
# NOTE deliberately NOT unignored: *.self_attn.kv_b_proj, *.self_attn.f_b_proj,
# *.self_attn.g_b_proj, *.self_attn.fused_fg_b_proj, indexer/conv/norm entries.


def read_header(path):
    with open(path, "rb") as f:
        n = struct.unpack("<Q", f.read(8))[0]
        return json.loads(f.read(n)), 8 + n


def block_for(name):
    """Return the quant block size for a target tensor name, else None."""
    m = MLP_RE.match(name)
    if m:
        layer, shared = int(m.group(1)), bool(m.group(2))
        return 128 if (shared or layer < FIRST_K_DENSE) else None
    if ATTN128_RE.match(name):
        return 128
    if ATTN32_RE.match(name):
        return 32
    return None


def quant_block(w32, b):
    n, k = w32.shape
    assert n % b == 0 and k % b == 0, (n, k, b)
    B = w32.reshape(n // b, b, k // b, b)
    amax = np.abs(B).max(axis=(1, 3), keepdims=True)
    s = np.where(amax == 0, np.float32(1), amax / np.float32(448)).astype(np.float32)
    q = np.clip(B / s, -448, 448).astype(ml_dtypes.float8_e4m3fn)
    return (
        np.ascontiguousarray(q.reshape(n, k)),
        np.ascontiguousarray(s.reshape(B.shape[0], B.shape[2])),
    )


def build_quantized_layers(mla_layers):
    ql = {}
    nvfp4 = {"quant_algo": "NVFP4", "group_size": 16}
    fp8_128 = {"quant_algo": "FP8_BLOCK128"}
    fp8_32 = {"quant_algo": "FP8_BLOCK32"}

    def attn_entries(base, layer):
        # base ends in ".self_attn"
        ql[f"{base}.o_proj"] = fp8_128
        if layer in mla_layers:
            for p in ("q_a_proj", "kv_a_proj_with_mqa", "q_b_proj"):
                ql[f"{base}.{p}"] = fp8_128
        else:
            for p in ("q_proj", "k_proj", "v_proj", "b_proj", "f_a_proj", "g_a_proj"):
                ql[f"{base}.{p}"] = fp8_32

    for layer in range(NUM_LAYERS):
        base = f"model.language_model.layers.{layer}"
        attn_entries(f"{base}.self_attn", layer)
        if layer >= FIRST_K_DENSE:
            for e in range(NUM_EXPERTS):
                for p in ("gate", "up", "down"):
                    ql[f"{base}.mlp.experts.{e}.{p}_proj"] = nvfp4
            for p in ("gate", "up", "down"):
                ql[f"{base}.mlp.shared_experts.{p}_proj"] = fp8_128
        else:
            for p in ("gate", "up", "down"):
                ql[f"{base}.mlp.{p}_proj"] = fp8_128

    # MTP draft-model runtime spelling: mtp.py builds the layer-45 decoder
    # block with prefix "model.layers.45" (no ".mtp_block" in the quant-lookup
    # prefix), so the manifest must carry these keys verbatim. Mirrors the MLP
    # aliases that fixed the FP8MIX-v2 KeyError; extended here to attention.
    mtp = "model.layers.45"
    attn_entries(f"{mtp}.self_attn", 45)
    for e in range(NUM_EXPERTS):
        for p in ("gate", "up", "down"):
            ql[f"{mtp}.mlp.experts.{e}.{p}_proj"] = nvfp4
    ql[f"{mtp}.mlp.experts.routed_experts.0.up_proj"] = nvfp4
    for p in ("gate", "up", "down"):
        ql[f"{mtp}.mlp.shared_experts.{p}_proj"] = fp8_128
    return ql


def main():
    os.makedirs(DST, exist_ok=True)
    idx = json.load(open(f"{SRC}/model.safetensors.index.json"))
    wmap = dict(idx["weight_map"])
    shards = sorted(set(wmap.values()))

    mla_layers = sorted(
        {
            int(re.match(r"model\.language_model\.layers\.(\d+)\.", n).group(1))
            for n in wmap
            if ".self_attn.q_a_proj." in n
        }
    )
    assert mla_layers == [3, 7, 11, 15, 19, 23, 27, 31, 35, 39, 43, 45], mla_layers

    todo = {}
    n32 = n128 = 0
    for name, sh in wmap.items():
        b = block_for(name)
        if b is not None:
            todo.setdefault(sh, {})[name] = b
            if b == 32:
                n32 += 1
            else:
                n128 += 1
    n_targets = n32 + n128
    print(f"{n_targets} targets ({n32} block32 + {n128} block128) in {len(todo)}/{len(shards)} shards")
    assert n32 == 204, n32          # 34 KDA layers x 6
    assert n128 == 138 + 82, n128   # MLP 138 + attn (46 o + 12x3 MLA)

    t0 = time.time()
    linked = rewritten = 0
    size_delta = 0
    for si, sh in enumerate(shards, 1):
        src, dst = f"{SRC}/{sh}", f"{DST}/{sh}"
        if os.path.exists(dst):
            os.unlink(dst)
        if sh not in todo:
            os.link(src, dst)
            linked += 1
            continue
        hdr, base = read_header(src)
        meta = hdr.get("__metadata__")
        conv = todo[sh]
        order = [k for k in hdr if k != "__metadata__"]
        new_hdr, blobs, off = {}, [], 0
        with open(src, "rb") as f:
            for k in order:
                m = hdr[k]
                s, e = m["data_offsets"]
                f.seek(base + s)
                raw = f.read(e - s)
                if k in conv:
                    assert m["dtype"] == "BF16", (k, m["dtype"])
                    w = (
                        np.frombuffer(raw, dtype=ml_dtypes.bfloat16)
                        .reshape(m["shape"])
                        .astype(np.float32)
                    )
                    qw, sc = quant_block(w, conv[k])
                    for nm, arr, dt in (
                        (k, qw, "F8_E4M3"),
                        (k[: -len(".weight")] + ".weight_scale_inv", sc, "F32"),
                    ):
                        b = arr.tobytes()
                        new_hdr[nm] = {
                            "dtype": dt,
                            "shape": list(arr.shape),
                            "data_offsets": [off, off + len(b)],
                        }
                        blobs.append(b)
                        off += len(b)
                        size_delta += len(b)
                    size_delta -= len(raw)
                else:
                    new_hdr[k] = {
                        "dtype": m["dtype"],
                        "shape": m["shape"],
                        "data_offsets": [off, off + len(raw)],
                    }
                    blobs.append(raw)
                    off += len(raw)
        if meta is not None:
            new_hdr["__metadata__"] = meta
        hb = json.dumps(new_hdr, separators=(",", ":")).encode()
        pad = (-(8 + len(hb))) % 8
        hb += b" " * pad
        tmp = dst + ".tmp"
        with open(tmp, "wb") as f:
            f.write(struct.pack("<Q", len(hb)))
            f.write(hb)
            for b in blobs:
                f.write(b)
        os.replace(tmp, dst)
        rewritten += 1
        print(
            f"[{si}/{len(shards)}] {sh} rewrote {len(conv)} tensors "
            f"({time.time() - t0:.0f}s)",
            flush=True,
        )

    for sh, names in todo.items():
        for name in names:
            wmap[name[: -len(".weight")] + ".weight_scale_inv"] = sh
    idx["weight_map"] = wmap
    if "metadata" in idx and "total_size" in idx["metadata"]:
        idx["metadata"]["total_size"] += size_delta
    json.dump(idx, open(f"{DST}/model.safetensors.index.json", "w"))

    cfg = json.load(open(f"{SRC}/config.json"))
    src_q = cfg["quantization_config"]
    ignore = [p for p in src_q["ignore"] if p not in UNIGNORE]
    cfg["quantization_config"] = {
        "producer": src_q.get("producer", {"name": "modelopt", "version": "0.45.0"}),
        "quant_method": "modelopt",
        "quant_algo": "MIXED_PRECISION",
        "kv_cache_scheme": src_q.get("kv_cache_scheme"),
        "group_size": 16,
        "ignore": ignore,
        "quantized_layers": build_quantized_layers(set(mla_layers)),
    }
    json.dump(cfg, open(f"{DST}/config.json", "w"), indent=1)

    for fn in os.listdir(SRC):
        if fn.endswith(".safetensors") or fn in (
            "config.json",
            "model.safetensors.index.json",
        ):
            continue
        s, d = f"{SRC}/{fn}", f"{DST}/{fn}"
        if os.path.isfile(s) and not os.path.exists(d):
            os.link(s, d)
    print(
        f"DONE linked={linked} rewritten={rewritten} "
        f"size_delta={size_delta / 2**30:.3f} GiB in {time.time() - t0:.0f}s"
    )


if __name__ == "__main__":
    import argparse

    ap = argparse.ArgumentParser(
        description="Derive the NVFP4 + block-FP8 (MLP + attention) checkpoint from a "
        "pristine GLM-5.3-Flash-NVFP4 checkpoint (rev caca4e6a). Unchanged shards are "
        "hardlinked (same filesystem required). Run convert_lmhead.py afterward to "
        "reproduce the published checkpoint exactly."
    )
    ap.add_argument(
        "--input",
        required=True,
        help="pristine GLM-5.3-Flash-NVFP4 checkpoint dir (only read)",
    )
    ap.add_argument(
        "--output",
        required=True,
        help="output checkpoint dir (created if missing; EXISTING FILES IN IT ARE REPLACED)",
    )
    ap.add_argument("--yes", action="store_true", help="skip the confirmation prompt")
    args = ap.parse_args()
    SRC, DST = os.path.abspath(args.input), os.path.abspath(args.output)
    if SRC == DST:
        raise SystemExit("--input and --output must differ (conversion is not in-place)")
    if not os.path.isfile(os.path.join(SRC, "model.safetensors.index.json")):
        raise SystemExit(f"{SRC} has no model.safetensors.index.json — not a checkpoint dir")
    if not args.yes:
        reply = input(f"(Re)write shards in {DST} from {SRC}? [y/N] ")
        if reply.strip().lower() not in ("y", "yes"):
            raise SystemExit("aborted")
    main()