#!/usr/bin/env python3
"""Assemble the hybrid NVFP4+AQLM GLM-5.2 checkpoint at /data/glm52.

Sources:
  - /tmp/glm52-dl/nvfp4-full: lukealonso/GLM-5.2-NVFP4 (all non-expert
    tensors verbatim; routed-expert tensors only for plan.nvfp4_layers)
  - /data/glm52-aqlm-parts/layer_N.pt: AQLM codes/codebooks/scales for
    the remaining expert layers

Output: sharded safetensors (~4GB each) + index + config.json with the
nvfp4_aqlm_hybrid quantization_config + tokenizer/aux files.
"""

import json
import os
import re
import shutil
import time

import torch
from safetensors import safe_open
from safetensors.torch import save_file

ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
SRC = "/tmp/glm52-dl/nvfp4-full"
PARTS = "/data/glm52-aqlm-parts"
DST = "/data/glm52"
SHARD_BYTES = 4 << 30

plan = json.load(open(os.path.join(ROOT, "hybrid_plan.json")))
NVFP4_LAYERS = set(plan["nvfp4_layers"])
AQLM_LAYERS = sorted(plan["aqlm_mixed_layers"]) + sorted(plan["aqlm_cold_layers"])


def keep(name: str) -> bool:
    m = re.match(r"model\.layers\.(\d+)\.mlp\.experts\.\d+\.", name)
    if m:
        return int(m.group(1)) in NVFP4_LAYERS
    return True


class ShardWriter:
    def __init__(self, dst: str):
        self.dst = dst
        self.cur: dict[str, torch.Tensor] = {}
        self.cur_bytes = 0
        self.n = 0
        self.weight_map: dict[str, str] = {}
        self.total = 0
        self.files: list[str] = []

    def add(self, name: str, tensor: torch.Tensor):
        nb = tensor.numel() * tensor.element_size()
        if self.cur_bytes + nb > SHARD_BYTES and self.cur:
            self.flush()
        self.cur[name] = tensor
        self.cur_bytes += nb
        self.total += nb

    def flush(self):
        if not self.cur:
            return
        self.n += 1
        fname = f"model-{self.n:05d}.safetensors"
        save_file(self.cur, os.path.join(self.dst, fname))
        for k in self.cur:
            self.weight_map[k] = fname
        print(f"  wrote {fname} ({self.cur_bytes/1e9:.2f} GB)", flush=True)
        self.files.append(fname)
        self.cur = {}
        self.cur_bytes = 0

    def finalize(self):
        self.flush()
        # rename to canonical model-XXXXX-of-NNNNN scheme
        total_n = self.n
        wm = {}
        for i, fname in enumerate(self.files, 1):
            new = f"model-{i:05d}-of-{total_n:05d}.safetensors"
            os.rename(
                os.path.join(self.dst, fname), os.path.join(self.dst, new)
            )
            for k, v in self.weight_map.items():
                if v == fname:
                    wm[k] = new
        idx = {
            "metadata": {"total_size": self.total},
            "weight_map": wm,
        }
        json.dump(
            idx,
            open(os.path.join(self.dst, "model.safetensors.index.json"), "w"),
            indent=0,
        )
        print(f"index: {len(wm)} tensors, {self.total/1e9:.1f} GB total")


def main():
    os.makedirs(DST, exist_ok=True)

    idx = json.load(open(os.path.join(SRC, "model.safetensors.index.json")))
    weight_map = idx["weight_map"]
    shards = sorted(set(weight_map.values()))

    writer = ShardWriter(DST)

    # 1. copy kept tensors from the NVFP4 repo, shard by shard
    for shard in shards:
        with safe_open(os.path.join(SRC, shard), framework="pt") as f:
            names = [n for n in f.keys() if keep(n)]
            if not names:
                continue
            print(f"{shard}: keeping {len(names)} tensors", flush=True)
            for n in names:
                writer.add(n, f.get_tensor(n))

    # 2. append AQLM expert tensors
    layer_books = {}
    mse_report = {}
    for li in AQLM_LAYERS:
        part = os.path.join(PARTS, f"layer_{li}.pt")
        # the quantizer writes in place: wait for existence + 60s of quiescence
        while (
            not os.path.exists(part) or time.time() - os.path.getmtime(part) < 60
        ):
            print(f"waiting for quantizer: layer {li} ...", flush=True)
            time.sleep(30)
        d = torch.load(part, map_location="cpu", weights_only=True)
        p = f"model.layers.{li}.mlp.experts"
        writer.add(f"{p}.w13_codes", d["w13_codes"])
        writer.add(f"{p}.w13_codebooks", d["w13_codebooks"])
        writer.add(f"{p}.w13_scales", d["w13_scales"])
        writer.add(f"{p}.w2_codes", d["w2_codes"])
        writer.add(f"{p}.w2_codebooks", d["w2_codebooks"])
        writer.add(f"{p}.w2_scales", d["w2_scales"])
        layer_books[str(li)] = {
            "w13": d["books"]["w13"],
            "w2": d["books"]["w2"],
        }
        mse_report[li] = (d["w13_rel_mse"], d["w2_rel_mse"])
        print(
            f"layer {li}: aqlm books={d['books']} "
            f"rel_mse w13={d['w13_rel_mse']:.4f} w2={d['w2_rel_mse']:.4f}",
            flush=True,
        )

    writer.finalize()

    # 3. config.json: replace quantization_config with the hybrid config,
    #    dropping quantized expert layers from the modelopt view is NOT
    #    needed -- AQLM layers are dispatched before the nvfp4 config is
    #    consulted, and their tensors do not match NVFP4 patterns.
    cfg = json.load(open(os.path.join(SRC, "config.json")))
    nvfp4_qc = cfg["quantization_config"]
    cfg["quantization_config"] = {
        "quant_method": "nvfp4_aqlm_hybrid",
        "nvfp4": nvfp4_qc,
        "aqlm": {"entries": 65536, "group_size": 8},
        "aqlm_layer_books": layer_books,
    }
    json.dump(cfg, open(os.path.join(DST, "config.json"), "w"), indent=2)

    # 4. aux files
    for f in os.listdir(SRC):
        if (
            f.endswith(".json")
            and f not in ("config.json", "model.safetensors.index.json")
            or f.endswith((".txt", ".jinja", ".py", ".md"))
        ):
            shutil.copy2(os.path.join(SRC, f), os.path.join(DST, f))
            print(f"copied {f}")

    json.dump(
        {str(k): v for k, v in mse_report.items()},
        open(os.path.join(DST, "aqlm_mse_report.json"), "w"),
        indent=2,
    )
    print("DONE:", DST)


if __name__ == "__main__":
    main()