#!/usr/bin/env python3
"""Patch /data/glm52 in place after demoting layers from NVFP4 to AQLM.

For each demoted layer: append its AQLM tensors as new shards, drop the
old per-expert NVFP4 tensor entries from the index (dead bytes stay in the
old shards, which is harmless), and update config.json aqlm_layer_books.
"""

import json
import os

import torch
from safetensors.torch import save_file

ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
PARTS = "/data/glm52-aqlm-parts"
DST = "/data/glm52"

plan = json.load(open(os.path.join(ROOT, "hybrid_plan.json")))
cfg = json.load(open(os.path.join(DST, "config.json")))
books_cfg = cfg["quantization_config"]["aqlm_layer_books"]

aqlm_all = set(plan["aqlm_mixed_layers"]) | set(plan["aqlm_cold_layers"])
demoted = sorted(li for li in aqlm_all if str(li) not in books_cfg)
print("layers to patch:", demoted)
assert demoted, "nothing to patch"

idx_path = os.path.join(DST, "model.safetensors.index.json")
idx = json.load(open(idx_path))
wm = idx["weight_map"]

for li in demoted:
    part = torch.load(
        os.path.join(PARTS, f"layer_{li}.pt"), map_location="cpu",
        weights_only=True,
    )
    prefix = f"model.layers.{li}.mlp.experts"
    tensors = {
        f"{prefix}.w13_codes": part["w13_codes"],
        f"{prefix}.w13_codebooks": part["w13_codebooks"],
        f"{prefix}.w13_scales": part["w13_scales"],
        f"{prefix}.w2_codes": part["w2_codes"],
        f"{prefix}.w2_codebooks": part["w2_codebooks"],
        f"{prefix}.w2_scales": part["w2_scales"],
    }
    fname = f"model-aqlm-patch-layer{li}.safetensors"
    save_file(tensors, os.path.join(DST, fname))

    # drop old per-expert NVFP4 entries for this layer
    stale = [n for n in wm if n.startswith(f"{prefix}.") and ".w13_" not in n
             and ".w2_" not in n]
    for n in stale:
        del wm[n]
    for n in tensors:
        wm[n] = fname
    books_cfg[str(li)] = {
        "w13": part["books"]["w13"], "w2": part["books"]["w2"],
    }
    nb = sum(t.numel() * t.element_size() for t in tensors.values())
    print(f"layer {li}: -{len(stale)} nvfp4 tensors, +{fname} "
          f"({nb/1e9:.2f} GB), rel_mse w13={part['w13_rel_mse']:.4f} "
          f"w2={part['w2_rel_mse']:.4f}")

# recompute logical total size
json.dump(idx, open(idx_path, "w"), indent=0)
json.dump(cfg, open(os.path.join(DST, "config.json"), "w"), indent=2)

# summarize logical size (sum over indexed tensors is expensive; report delta)
print("index now has", len(wm), "tensors")
print("aqlm_layer_books entries:", len(books_cfg))
print("DONE")