"""Post-step on GLM-5.3-Flash-NVFP4-FP8ATTN: quantize lm_head to block-FP8.

Composable AFTER convert_fp8attn.py (r2 content). Targets the fp8attn-r4
image, whose modelopt.py dispatches ParallelLMHead + FP8_BLOCK128 to
_Fp8BlockLMHeadMethod (vocab-block-sharded scale loader).

  lm_head.weight  BF16 [154880, 4096]  ->  F8_E4M3 + weight_scale_inv F32
  [1210, 32], block [128,128], dequant multiplier amax/448 — identical
  scheme/derivation to convert_fp8attn.py's quant_block.

Geometry: 154880 % 128 == 0 and the TP2 vocab shard 77440 % 128 == 0, so
weight and scale shard exactly on block boundaries.

config.json: remove "lm_head" from quantization_config.ignore (its substring
match would keep the head excluded) and add
quantized_layers["lm_head"] = {"quant_algo": "FP8_BLOCK128"}.

Reversible: the touched shard, config.json and the index are snapshotted as
*.pre-lmhead first (shard via hardlink — zero cost). `--revert` puts all
three back and removes the snapshots. embed_tokens stays BF16 (gather
semantics: VocabParallelEmbedding never receives a quant method).

Source NVFP4 dir untouched (all writes are new-file + os.replace).
"""

import json
import os
import struct
import sys
import time

import ml_dtypes
import numpy as np

DST = None  # set from --dir in __main__
TENSOR = "lm_head.weight"
BLOCK = 128
SUFFIX = ".pre-lmhead"


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 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 revert():
    idx_path = f"{DST}/model.safetensors.index.json"
    cfg_path = f"{DST}/config.json"
    idx_bak, cfg_bak = idx_path + SUFFIX, cfg_path + SUFFIX
    if not os.path.exists(idx_bak):
        sys.exit("no .pre-lmhead snapshots found; nothing to revert")
    shard = json.load(open(idx_bak))["weight_map"][TENSOR]
    shard_path = f"{DST}/{shard}"
    for bak, live in ((shard_path + SUFFIX, shard_path), (idx_bak, idx_path), (cfg_bak, cfg_path)):
        os.replace(bak, live)
    print(f"reverted: {shard}, index, config restored")


def main():
    idx_path = f"{DST}/model.safetensors.index.json"
    cfg_path = f"{DST}/config.json"
    idx = json.load(open(idx_path))
    if TENSOR + "_scale_inv" in idx["weight_map"] or (
        TENSOR[: -len(".weight")] + ".weight_scale_inv"
    ) in idx["weight_map"]:
        sys.exit("lm_head already converted (scale in index); use --revert first")
    shard = idx["weight_map"][TENSOR]
    src = f"{DST}/{shard}"

    # snapshots (shard as hardlink: pre-conversion inode preserved for free)
    os.link(src, src + SUFFIX)
    for p in (idx_path, cfg_path):
        with open(p) as f, open(p + SUFFIX, "w") as g:
            g.write(f.read())

    t0 = time.time()
    hdr, base = read_header(src)
    meta = hdr.get("__metadata__")
    order = [k for k in hdr if k != "__metadata__"]
    new_hdr, blobs, off, size_delta = {}, [], 0, 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 == TENSOR:
                assert m["dtype"] == "BF16", m["dtype"]
                w = (
                    np.frombuffer(raw, dtype=ml_dtypes.bfloat16)
                    .reshape(m["shape"])
                    .astype(np.float32)
                )
                qw, sc = quant_block(w, BLOCK)
                # round-trip check
                deq = qw.astype(np.float32) * np.repeat(
                    np.repeat(sc, BLOCK, 0), BLOCK, 1
                )
                err = np.linalg.norm(deq - w) / np.linalg.norm(w)
                print(f"lm_head shape {m['shape']} round-trip rel err {err:.4%}")
                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()
    hb += b" " * ((-(8 + len(hb))) % 8)
    tmp = src + ".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, src)

    wmap = dict(idx["weight_map"])
    wmap[TENSOR[: -len(".weight")] + ".weight_scale_inv"] = shard
    idx["weight_map"] = wmap
    if "metadata" in idx and "total_size" in idx["metadata"]:
        idx["metadata"]["total_size"] += size_delta
    json.dump(idx, open(idx_path, "w"))

    cfg = json.load(open(cfg_path))
    q = cfg["quantization_config"]
    q["ignore"] = [p for p in q["ignore"] if p != "lm_head"]
    q["quantized_layers"]["lm_head"] = {"quant_algo": "FP8_BLOCK128"}
    json.dump(cfg, open(cfg_path, "w"), indent=1)

    print(
        f"DONE {shard} rewritten, size_delta={size_delta / 2**30:+.3f} GiB, "
        f"{time.time() - t0:.0f}s. Revert: python3 convert_lmhead.py --dir {DST} --revert"
    )


if __name__ == "__main__":
    import argparse

    ap = argparse.ArgumentParser(
        description="Quantize lm_head to block-FP8 IN PLACE in a converted "
        "GLM-5.3-Flash-NVFP4-FP8ATTN checkpoint dir. The touched shard, config.json and "
        "index are snapshotted as *.pre-lmhead first; --revert restores them."
    )
    ap.add_argument(
        "--dir", required=True, help="checkpoint dir to modify IN PLACE"
    )
    ap.add_argument(
        "--revert", action="store_true", help="restore the *.pre-lmhead snapshots"
    )
    ap.add_argument("--yes", action="store_true", help="skip the confirmation prompt")
    args = ap.parse_args()
    DST = os.path.abspath(args.dir)
    if not os.path.isfile(os.path.join(DST, "model.safetensors.index.json")):
        raise SystemExit(f"{DST} has no model.safetensors.index.json — not a checkpoint dir")
    if args.revert:
        revert()
    else:
        if not args.yes:
            reply = input(
                f"Rewrite the lm_head shard in {DST} IN PLACE "
                "(with .pre-lmhead snapshots)? [y/N] "
            )
            if reply.strip().lower() not in ("y", "yes"):
                raise SystemExit("aborted")
        main()