"""Derive a block-FP8-visual-tower checkpoint from GLM-5.3-Flash-NVFP4-FP8ATTN-recal.

Worker VISION (glm-max push). The r2-era "yields NaN image features" warning in
model.py (~line 1050) is about inheriting the global fp8 quant_config while the
checkpoint ships BF16 visual weights WITHOUT weight_scale_inv — i.e. a
loader/format mismatch, not FP8 numerics of the tower. This conversion ships
the scales, so with the manifest-gated model.py overlay (vision-test image) the
tower's Linear layers load through vLLM's proven Fp8LinearMethod block path
(same serialization as the r2 MLP/attention conversion: weight F8_E4M3 [N,K] +
weight_scale_inv F32 [N/128,K/128], dequant multiplier amax/448).

Converted (124 tensors, 1.0156 GiB BF16 -> ~0.512 GiB, all in shard 120):
  model.visual.blocks.{0..23}.attn.qkv.weight        [3072,1024]
  model.visual.blocks.{0..23}.attn.proj.weight       [1024,1024]
  model.visual.blocks.{0..23}.mlp.{gate,up,down}_proj.weight
  model.visual.merger.{proj,gate_proj,up_proj,down_proj}.weight
NOT converted: patch_embed (Conv3d), downsample (Conv2d), all biases/norms
(34.2 MiB total). Every shape and TP2 shard offset divides by 128 (verified:
qkv q/k/v sections 1024 each -> 512/rank; merger 10240 -> 5120/rank).

Config: prune the three visual ignore patterns; add FP8_BLOCK128
quantized_layers entries under the RUNTIME lookup prefixes
("visual.blocks.N.attn.qkv_proj" — note qkv_proj, multimodal.py renames the
prefix when quant_config is passed) plus "model.visual." spellings and
per-shard aliases (strategy-4 fused lookups), mirroring FIT's r4 alias
precedent. Unlisted visual linears resolve to UnquantizedLinearMethod, so
pruning the ignore patterns is safe for the conv/norm remainder.

Usage: python3 convert_fp8visual.py            # recal -> -recal-visfp8
       SRC=<dir> DST=<dir> python3 ...         # stack on DRAFT's dir
"""

import json
import os
import re
import struct
import time

import ml_dtypes
import numpy as np

SRC = os.environ.get("SRC", "/home/ian/models/GLM-5.3-Flash-NVFP4-FP8ATTN-recal")
DST = os.environ.get(
    "DST", "/home/ian/models/GLM-5.3-Flash-NVFP4-FP8ATTN-recal-visfp8"
)
BLOCK = 128
DEPTH = 24

TARGET_RE = re.compile(
    r"^model\.visual\.("
    r"blocks\.\d+\.attn\.(qkv|proj)"
    r"|blocks\.\d+\.mlp\.(gate|up|down)_proj"
    r"|merger\.(proj|gate_proj|up_proj|down_proj)"
    r")\.weight$"
)

VISUAL_IGNORE = {"model.visual.*", "visual.*", "*.visual.*"}


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=BLOCK):
    n, k = w32.shape
    assert n % b == 0 and k % b == 0, (n, k)
    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 visual_manifest_entries():
    fp8 = {"quant_algo": "FP8_BLOCK128"}
    mods = []
    for i in range(DEPTH):
        mods += [
            # runtime lookup prefix (multimodal.py: ".qkv_proj" iff quant_config)
            f"visual.blocks.{i}.attn.qkv_proj",
            f"visual.blocks.{i}.attn.qkv",
            f"visual.blocks.{i}.attn.proj",
            f"visual.blocks.{i}.mlp.gate_up_proj",
            f"visual.blocks.{i}.mlp.gate_proj",
            f"visual.blocks.{i}.mlp.up_proj",
            f"visual.blocks.{i}.mlp.down_proj",
        ]
    mods += [
        "visual.merger.proj",
        "visual.merger.gate_up_proj",
        "visual.merger.gate_proj",
        "visual.merger.up_proj",
        "visual.merger.down_proj",
    ]
    ql = {}
    for m in mods:
        ql[m] = fp8
        ql["model." + m] = fp8  # in case apply_vllm_mapper is/isn't applied
    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()))

    todo = {}
    for name, sh in wmap.items():
        if TARGET_RE.match(name):
            todo.setdefault(sh, set()).add(name)
    n_targets = sum(len(v) for v in todo.values())
    print(f"{n_targets} target tensors in {len(todo)}/{len(shards)} shards")
    assert n_targets == 124, n_targets

    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)
                    # round-trip check
                    deq = qw.astype(np.float32) * np.kron(
                        sc, np.ones((BLOCK, BLOCK), np.float32)
                    )
                    rel = np.linalg.norm(deq - w) / np.linalg.norm(w)
                    if rel > 0.05:
                        print(f"  WARN {k} round-trip rel err {rel:.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()
        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"))
    q = cfg["quantization_config"]
    q["ignore"] = [p for p in q["ignore"] if p not in VISUAL_IGNORE]
    q["quantized_layers"] = {**q["quantized_layers"], **visual_manifest_entries()}
    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__":
    main()