#!/usr/bin/env python3
"""AQLM quantization of GLM-5.2 routed experts from the NVFP4 checkpoint.

For each AQLM-assigned layer (see hybrid_plan.json):
  1. Load the layer's 256 experts' NVFP4 tensors from the downloaded repo,
     dequantize to bf16 on GPU (u8-packed fp4 * fp8 block scale * fp32
     global scale).
  2. Per projection (w13 = gate rows then up rows; w2 = down):
     a. per-output-channel scale s[e,o] (mean abs of the row);
     b. fit one 65536 x 8 codebook per book with GPU k-means on a row
        sample (codebook shared across all 256 experts);
     c. encode all groups-of-8 by chunked nearest-centroid search
        (books=2 encodes the residual against a second codebook);
     d. least-squares refit of the per-channel scales.
  3. Save codes/codebooks/scales + reconstruction stats per layer.

Everything is chunked: peak GPU memory is the bf16 expert tensor plus a
few GB of scratch, so 8 workers (one per GPU) run comfortably.

Usage: aqlm_quantize.py [--layers 11,12] [--gpus 0,1,...] [--out DIR]
"""

import argparse
import json
import os
import time
from concurrent.futures import ProcessPoolExecutor

ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
NVFP4_DIR = os.environ.get("GLM52_NVFP4_DIR", "/tmp/glm52-dl/nvfp4-full")
OUT_DIR = "/data/glm52-aqlm-parts"
ENTRIES = 65536
GDIM = 8
SAMPLE_VECS = 4_000_000
KMEANS_ITERS = 12
CHUNK_VECS = 32_768  # scores are [CHUNK_VECS, 65536] fp16 (~4.3GB)

FP4_LUT = [0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0,
           -0.0, -0.5, -1.0, -1.5, -2.0, -3.0, -4.0, -6.0]


def load_plan():
    plan = json.load(open(os.path.join(ROOT, "hybrid_plan.json")))

    def books(tier):
        t = plan["aqlm"][tier]
        return {"w13": t["w13_books"], "w2": t["w2_books"]}

    layers = {}
    for li in plan["aqlm_mixed_layers"]:
        layers[li] = books("mixed")
    for li in plan["aqlm_cold_layers"]:
        layers[li] = books("cold")
    return layers


class ShardReader:
    """Random access to tensors across the downloaded safetensors shards."""

    def __init__(self, repo_dir: str):
        idx = json.load(
            open(os.path.join(repo_dir, "model.safetensors.index.json"))
        )
        self.weight_map = idx["weight_map"]
        self.repo_dir = repo_dir
        self._open = {}

    def get(self, name: str):
        from safetensors import safe_open

        shard = self.weight_map[name]
        if shard not in self._open:
            self._open[shard] = safe_open(
                os.path.join(self.repo_dir, shard), framework="pt"
            )
        return self._open[shard].get_tensor(name)


def dequant_nvfp4(reader, prefix: str, device, lut) -> "torch.Tensor":
    """Dequantize one NVFP4 linear weight to bf16 [out, in]."""
    import torch

    packed = reader.get(f"{prefix}.weight").to(device)          # u8 [out, in/2]
    wscale = reader.get(f"{prefix}.weight_scale").to(device)    # fp8 [out, in/16]
    wscale2 = reader.get(f"{prefix}.weight_scale_2").to(device)  # fp32 scalar

    lo = (packed & 0x0F).long()
    hi = (packed >> 4).long()
    vals = torch.empty(
        packed.shape[0], packed.shape[1] * 2, dtype=torch.float32, device=device
    )
    # NVFP4 packs the first element of each pair in the low nibble.
    vals[:, 0::2] = lut[lo]
    vals[:, 1::2] = lut[hi]
    del lo, hi
    scale = wscale.to(torch.float32).repeat_interleave(16, dim=1)
    return (vals * scale * wscale2.to(torch.float32)).to(torch.bfloat16)


def _argmin_codes(v16, cent16, cn16):
    """v16 [n,8] fp16, cent16 [ENTRIES,8] fp16 -> nearest ids int32.
    argmax of 2*v.c - |c|^2 (|v|^2 constant per vector)."""
    import torch

    d = v16 @ cent16.t()
    d.mul_(2).sub_(cn16)
    return d.argmax(-1).to(torch.int32)


def encode_chunked(vecs16, cent16):
    import torch

    cn16 = (cent16.float() * cent16.float()).sum(-1).half()
    out = torch.empty(
        vecs16.shape[0], dtype=torch.int32, device=vecs16.device
    )
    for s in range(0, vecs16.shape[0], CHUNK_VECS):
        out[s : s + CHUNK_VECS] = _argmin_codes(
            vecs16[s : s + CHUNK_VECS], cent16, cn16
        )
    return out


def kmeans_fit(vecs16, iters: int, generator):
    """Lloyd's k-means on [N,8] fp16; returns [ENTRIES,8] fp16 centroids."""
    import torch

    n = vecs16.shape[0]
    device = vecs16.device
    perm = torch.randperm(n, generator=generator, device=device)[:ENTRIES]
    cent = vecs16[perm].float()
    for _ in range(iters):
        cent16 = cent.half()
        cn16 = (cent * cent).sum(-1).half()
        sums = torch.zeros(ENTRIES, GDIM, device=device)
        cnts = torch.zeros(ENTRIES, device=device)
        ones = None
        for s in range(0, n, CHUNK_VECS):
            v = vecs16[s : s + CHUNK_VECS]
            a = _argmin_codes(v, cent16, cn16).long()
            sums.index_add_(0, a, v.float())
            if ones is None or ones.shape[0] != a.shape[0]:
                ones = torch.ones(a.shape[0], device=device)
            cnts.index_add_(0, a, ones[: a.shape[0]])
        mask = cnts > 0
        cent[mask] = sums[mask] / cnts[mask].unsqueeze(-1)
        dead = int((~mask).sum())
        if dead:
            ridx = torch.randperm(n, generator=generator, device=device)[:dead]
            cent[~mask] = vecs16[ridx].float()
    return cent.half()


def quantize_matrix(w, books: int, generator, log):
    """w: [E, out, in] bf16 (GPU) -> dict with codes int16 [E,books,out,in/8],
    codebooks fp16 [books,ENTRIES,8], scales fp16 [E,out], rel mse."""
    import torch

    e, out, k = w.shape
    device = w.device
    rows = e * out
    k8 = k // GDIM
    wv = w.reshape(rows, k)

    # per-row scales, chunked
    scales = torch.empty(rows, dtype=torch.float32, device=device)
    ROWCH = max(1, (CHUNK_VECS * 64) // k)
    for s in range(0, rows, ROWCH):
        scales[s : s + ROWCH] = (
            wv[s : s + ROWCH].float().abs().mean(-1).clamp_min(1e-8)
        )

    # --- fit codebooks on a row sample ---
    n_sample_rows = min(rows, max(1, SAMPLE_VECS // k8))
    ridx = torch.randperm(rows, generator=generator, device=device)[
        :n_sample_rows
    ]
    samp = (
        (wv[ridx].float() / scales[ridx].unsqueeze(-1))
        .reshape(-1, GDIM)
        .half()
    )
    log(f"    kmeans book0 on {samp.shape[0]} vecs")
    cbs = [kmeans_fit(samp, KMEANS_ITERS, generator)]
    if books == 2:
        c0 = encode_chunked(samp, cbs[0])
        resid = (samp.float() - cbs[0].float()[c0.long()]).half()
        log(f"    kmeans book1 on residuals")
        cbs.append(kmeans_fit(resid, KMEANS_ITERS, generator))
        del c0, resid
    del samp
    torch.cuda.empty_cache()

    cb_f = [c.float() for c in cbs]
    cn16 = [(cf * cf).sum(-1).half() for cf in cb_f]

    # --- full encode, chunked by rows; accumulate scale-refit stats ---
    codes = torch.empty(
        (rows, books, k8), dtype=torch.int16, device=device
    )
    num = torch.zeros(rows, dtype=torch.float32, device=device)
    den = torch.zeros(rows, dtype=torch.float32, device=device)
    sse = 0.0
    stot = 0.0
    ROWCH = max(1, CHUNK_VECS // k8)
    for s in range(0, rows, ROWCH):
        wc = wv[s : s + ROWCH].float() / scales[s : s + ROWCH].unsqueeze(-1)
        v = wc.reshape(-1, GDIM).half()
        approx = torch.zeros(v.shape[0], GDIM, device=device)
        r = v.float()
        for b in range(books):
            cb = _argmin_codes(r.half(), cbs[b], cn16[b]).long()
            dec = cb_f[b][cb]
            approx += dec
            r -= dec
            codes[s : s + ROWCH, b] = (
                torch.where(cb >= 32768, cb - 65536, cb)
                .to(torch.int16)
                .reshape(-1, k8)
            )
        nrow = wc.shape[0]
        a2 = approx.reshape(nrow, k)
        w2 = wc
        num[s : s + ROWCH] = (w2 * a2).sum(-1)
        den[s : s + ROWCH] = (a2 * a2).sum(-1).clamp_min(1e-8)
        sse += float((r * r).sum())
        stot += float((w2 * w2).sum())
        del wc, v, approx, r, a2

    rel = sse / max(stot, 1e-9)
    scales = scales * (num / den).clamp(0.5, 2.0)
    log(f"    rel_mse={rel:.4f}")

    return {
        "codes": codes.reshape(e, out, books, k8)
        .permute(0, 2, 1, 3)
        .contiguous()
        .cpu(),
        "codebooks": torch.stack(cbs).cpu(),
        "scales": scales.reshape(e, out).half().cpu(),
        "rel_mse": rel,
    }


def process_layer(layer_idx: int, books: dict, gpu: int):
    import torch

    t0 = time.time()
    device = f"cuda:{gpu}"
    torch.cuda.set_device(gpu)
    gen = torch.Generator(device=device)
    gen.manual_seed(1234 + layer_idx)
    lut = torch.tensor(FP4_LUT, dtype=torch.float32, device=device)

    def log(msg):
        print(f"[L{layer_idx} gpu{gpu}] {msg}", flush=True)

    out_path = os.path.join(OUT_DIR, f"layer_{layer_idx}.pt")
    if os.path.exists(out_path):
        log("already done, skip")
        return layer_idx, None

    reader = ShardReader(NVFP4_DIR)
    n_exp, inter, hidden = 256, 2048, 6144

    log("dequantizing experts")
    w13 = torch.empty(
        n_exp, 2 * inter, hidden, dtype=torch.bfloat16, device=device
    )
    w2 = torch.empty(n_exp, hidden, inter, dtype=torch.bfloat16, device=device)
    for ei in range(n_exp):
        p = f"model.layers.{layer_idx}.mlp.experts.{ei}"
        w13[ei, :inter] = dequant_nvfp4(reader, f"{p}.gate_proj", device, lut)
        w13[ei, inter:] = dequant_nvfp4(reader, f"{p}.up_proj", device, lut)
        w2[ei] = dequant_nvfp4(reader, f"{p}.down_proj", device, lut)

    log(f"w13 {tuple(w13.shape)} books={books['w13']}")
    w13_out = quantize_matrix(w13, books["w13"], gen, log)
    del w13
    torch.cuda.empty_cache()
    log(f"w2 {tuple(w2.shape)} books={books['w2']}")
    w2_out = quantize_matrix(w2, books["w2"], gen, log)
    del w2
    torch.cuda.empty_cache()

    torch.save(
        {
            "layer": layer_idx,
            "books": books,
            "w13_codes": w13_out["codes"],
            "w13_codebooks": w13_out["codebooks"],
            "w13_scales": w13_out["scales"],
            "w13_rel_mse": w13_out["rel_mse"],
            "w2_codes": w2_out["codes"],
            "w2_codebooks": w2_out["codebooks"],
            "w2_scales": w2_out["scales"],
            "w2_rel_mse": w2_out["rel_mse"],
        },
        out_path,
    )
    log(f"saved ({time.time()-t0:.0f}s)")
    return layer_idx, (w13_out["rel_mse"], w2_out["rel_mse"])


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--layers", default=None)
    ap.add_argument("--gpus", default="0,1,2,3,4,5,6,7")
    ap.add_argument("--out", default=OUT_DIR)
    args = ap.parse_args()

    globals()["OUT_DIR"] = args.out
    os.makedirs(args.out, exist_ok=True)
    plan = load_plan()
    layers = sorted(plan)
    if args.layers:
        want = {int(x) for x in args.layers.split(",")}
        layers = [li for li in layers if li in want]
    gpus = [int(g) for g in args.gpus.split(",")]

    print(f"{len(layers)} layers over {len(gpus)} GPUs", flush=True)
    jobs = [(li, plan[li], gpus[i % len(gpus)]) for i, li in enumerate(layers)]
    with ProcessPoolExecutor(max_workers=len(gpus)) as ex:
        futs = [ex.submit(process_layer, *j) for j in jobs]
        for f in futs:
            li, rel = f.result()
            if rel:
                print(
                    f"== layer {li}: rel_mse w13={rel[0]:.4f} w2={rel[1]:.4f}",
                    flush=True,
                )


if __name__ == "__main__":
    main()