"""Sarvam-Prism -> INT4 package for <=16 GB GPUs: routed experts to INT4 (GPTQ, group 128, symmetric), everything else kept as shipped.

Stages (one process; calibration activations live in RAM):
  1. load the original FP8 model on the big GPU, run the evaluation suite -> ref.json
  2. run calibration requests, record each MoE layer's input tokens (and their routing)
  3. free the GPU; for each MoE layer: build per-expert Hessians, GPTQ all 128 experts batched on the GPU,
     pack to int4, write the output layer file, delete the source layer file (space on /dev/shm is limited)
"""
import gc, json, os, shutil, sys, time
import torch
import torch.nn.functional as F
from safetensors import safe_open
from safetensors.torch import save_file

SRC, OUT, WORK = "/dev/shm/src", "/dev/shm/out", "/workspace/work"
GROUP, QMAX, QMIN = 128, 7, -8
CALIB_TOKENS = int(os.environ.get("CALIB_TOKENS", 262144))
CLIPS = torch.tensor([1.0, 0.96, 0.92, 0.88, 0.84, 0.80])
os.makedirs(OUT, exist_ok=True); os.makedirs(WORK, exist_ok=True)
log = lambda *a: print(time.strftime("%H:%M:%S"), *a, flush=True)


# ------------------------------------------------------------------------------------------------ stage 1 + 2
def reference_and_calibrate():
    sys.path.insert(0, SRC); sys.path.insert(0, os.path.dirname(__file__)); sys.path.insert(0, WORK)
    import prism
    from suite import run
    from calib_requests import make_requests
    t = time.time()
    model = prism.Prism.from_pretrained(SRC, cuda_graphs=False)
    log(f"reference loaded in {time.time()-t:.0f}s, GPU mem {torch.cuda.memory_allocated()/1e9:.1f} GB")
    if not os.path.exists(f"{WORK}/ref.json"):
        run(model, f"{WORK}/ref.json"); log("reference suite done")
    store, state = {}, {"on": False}
    for name, m in model.backbone.named_modules():
        if type(m).__name__ == "SarvamMLAMoE":
            li = int(name.split(".")[1]); orig = m._moe_dispatch; store[li] = ([], [])
            def wrap(x, ids, w, li=li, orig=orig):
                if state["on"]:
                    lens = prism.SPLIT["len"]; B = len(lens); L = x.shape[0] // B
                    sel = torch.cat([torch.arange(b * L, b * L + n, device=x.device) for b, n in enumerate(lens)])
                    store[li][0].append(x[sel].to(torch.bfloat16).cpu()); store[li][1].append(ids[sel].to(torch.uint8).cpu())
                return orig(x, ids, w)
            m._moe_dispatch = wrap
    reqs = make_requests(os.path.join(os.path.dirname(__file__), "wiki.json"))
    state["on"] = True; n_tok, t = 0, time.time()
    for i, (s, q) in enumerate(reqs):
        model.decide(state=s, questions=q)
        n_tok = sum(x.shape[0] for x in store[1][0])
        if i % 200 == 0: log(f"calib req {i} tokens {n_tok}")
        if n_tok >= CALIB_TOKENS: break
    state["on"] = False
    log(f"calibration: {i+1} requests, {n_tok} tokens, {time.time()-t:.0f}s")
    calib = {li: (torch.cat(xs), torch.cat(ids)) for li, (xs, ids) in store.items()}
    del model, store; gc.collect(); torch.cuda.empty_cache()
    return calib


# ------------------------------------------------------------------------------------------------ GPTQ
def group_scale(w):
    """w: [E, N, G] -> per-(E,N) scale chosen from CLIPS by minimum rounding MSE."""
    amax = w.abs().amax(-1).clamp_min(1e-8)
    best, best_err = None, None
    for c in CLIPS.tolist():
        s = amax * c / QMAX
        err = ((torch.clamp(torch.round(w / s[..., None]), QMIN, QMAX) * s[..., None] - w) ** 2).sum(-1)
        if best is None:
            best, best_err = s, err
        else:
            m = err < best_err; best = torch.where(m, s, best); best_err = torch.where(m, err, best_err)
    return best


@torch.no_grad()
def gptq(W, H, cnt, damp=0.01, min_tokens=512):
    """Batched GPTQ. W: [E, N, K] fp32, H: [E, K, K] fp32 (sum of x x^T), cnt: calibration tokens per expert.
    Experts seen by fewer than min_tokens tokens get H = I (plain round-to-nearest with clip search: no error
    feedback fitted to a handful of samples). Inputs never seen (zero diagonal) keep their weights (diag set to the
    mean) so rare inputs at deployment are still served. Returns int q [E,N,K] and fp32 scales [E,N,K/G]."""
    E, N, K = W.shape
    W = W.clone(); H = H.clone()
    few = (cnt < min_tokens).to(H.device)
    if few.any():
        H[few] = torch.eye(K, device=H.device)
    diag = torch.diagonal(H, dim1=1, dim2=2)
    mean = diag.mean(-1, keepdim=True).clamp_min(1e-8)
    dead = diag <= 0
    diag.copy_(torch.where(dead, mean.expand_as(diag), diag))
    diag += damp * mean
    L, info = torch.linalg.cholesky_ex(H)
    if info.any():                                            # extra damping for ill-conditioned experts
        bad = info.nonzero().flatten()
        for e in bad.tolist():
            d = torch.diagonal(H[e]); d += 0.1 * d.mean()
        L, info = torch.linalg.cholesky_ex(H)
        assert not info.any(), "Cholesky failed after extra damping"
    Hinv = torch.cholesky_inverse(L)
    Hinv = torch.linalg.cholesky(Hinv, upper=True)
    Q = torch.empty(E, N, K, dtype=torch.int8, device=W.device)
    S = torch.empty(E, N, K // GROUP, dtype=torch.float32, device=W.device)
    for g, i1 in enumerate(range(0, K, GROUP)):
        i2 = i1 + GROUP
        W1 = W[:, :, i1:i2].clone(); Err = torch.empty_like(W1); Hi = Hinv[:, i1:i2, i1:i2]
        s = group_scale(W1); S[:, :, g] = s
        for i in range(GROUP):
            w = W1[:, :, i]
            q = torch.clamp(torch.round(w / s), QMIN, QMAX)
            Q[:, :, i1 + i] = q.to(torch.int8)
            err = (w - q * s) / Hi[:, i, i][:, None]
            W1[:, :, i:] -= err[:, :, None] * Hi[:, i, i:][:, None, :]
            Err[:, :, i] = err
        if i2 < K:
            W[:, :, i2:] -= torch.bmm(Err, Hinv[:, i1:i2, i2:])
    return Q, S


def rtn(W):
    E, N, K = W.shape
    Wg = W.view(E, N, K // GROUP, GROUP)
    s = Wg.abs().amax(-1).clamp_min(1e-8) / QMAX
    return torch.clamp(torch.round(Wg / s[..., None]), QMIN, QMAX).view(E, N, K).to(torch.int8), s


def dequant(Q, S):
    E, N, K = Q.shape
    return (Q.float().view(E, N, K // GROUP, GROUP) * S[..., None]).view(E, N, K)


def pack(Q):
    """int q in [-8,7], [E,N,K] -> uint8 [E,N,K/2]. Within each 128-group, byte j holds element j (low nibble) and
    element j+64 (high nibble)."""
    E, N, K = Q.shape
    u = (Q.to(torch.int16) + 8).to(torch.uint8).view(E, N, K // GROUP, 2, GROUP // 2)
    return (u[:, :, :, 0] | (u[:, :, :, 1] << 4)).reshape(E, N, K // 2).contiguous()


def unpack(P):
    E, N, K2 = P.shape
    p = P.view(E, N, K2 // (GROUP // 2), GROUP // 2)
    lo, hi = (p & 15).to(torch.int8) - 8, (p >> 4).to(torch.int8) - 8
    return torch.stack([lo, hi], 3).reshape(E, N, K2 * 2)


def hessians(X, ids, n_exp, fn=None, dim=None):
    """Per-expert sum x x^T over the tokens routed to it. fn(e, x) maps the token inputs (e.g. to the down-proj
    input). Returns H [E, D, D] fp32 and token counts."""
    flat = ids.long()
    D = dim or X.shape[1]
    H = torch.zeros(n_exp, D, D, device="cuda")
    cnt = torch.zeros(n_exp, dtype=torch.long)
    tok = torch.arange(X.shape[0], device=X.device).repeat_interleave(ids.shape[1])
    order = flat.flatten().argsort()
    e_sorted, t_sorted = flat.flatten()[order], tok[order]
    bounds = torch.searchsorted(e_sorted, torch.arange(n_exp + 1, device=X.device)).tolist()
    for e in range(n_exp):
        t = t_sorted[bounds[e]:bounds[e + 1]]
        cnt[e] = t.numel()
        for c in range(0, t.numel(), 32768):
            x = X[t[c:c + 32768]].float()
            if fn is not None: x = fn(e, x)
            H[e].addmm_(x.T, x)
    return H, cnt


def rel_err(Wq, W, H):
    """Hessian-weighted relative output error: tr(D H D^T) / tr(W H W^T), averaged over experts."""
    D = Wq - W
    num = torch.einsum("enk,ekj,enj->e", D, H, D); den = torch.einsum("enk,ekj,enj->e", W, H, W)
    return (num / den.clamp_min(1e-12)).mean().item()


# ------------------------------------------------------------------------------------------------ stage 3
def quantize_layer(li, X, ids, stats):
    f = f"model-layer-{li:02d}.safetensors"
    tensors = {}
    with safe_open(f"{SRC}/{f}", "pt", device="cuda") as sf:
        for k in sf.keys(): tensors[k] = sf.get_tensor(k)
    p = f"model.layers.{li}.mlp.experts."
    gu = tensors.pop(p + "gate_up_proj").float() * tensors.pop(p + "gate_up_proj_scale")[..., None]   # [E, 2I, H]
    dn = tensors.pop(p + "down_proj").float() * tensors.pop(p + "down_proj_scale")[..., None]         # [E, H, I]
    E, I2, Hd = gu.shape; I = I2 // 2
    X, ids = X.cuda(), ids.cuda()
    t = time.time()
    Hgu, cnt = hessians(X, ids, E)
    Qgu, Sgu = gptq(gu, Hgu, cnt)
    guq = dequant(Qgu, Sgu)
    e_gu = rel_err(guq, gu, Hgu); e_gu_rtn = rel_err(dequant(*rtn(gu)), gu, Hgu); del Hgu
    def act(e, x):                                           # down-proj input through the *quantized* gate/up
        h = x @ guq[e].T
        return F.silu(h[:, :I]) * h[:, I:]
    Hdn, _ = hessians(X, ids, E, fn=act, dim=I)
    Qdn, Sdn = gptq(dn, Hdn, cnt)
    e_dn = rel_err(dequant(Qdn, Sdn), dn, Hdn); e_dn_rtn = rel_err(dequant(*rtn(dn)), dn, Hdn); del Hdn
    assert torch.equal(unpack(pack(Qgu)), Qgu) and torch.equal(unpack(pack(Qdn)), Qdn)
    tensors[p + "gate_up_qweight"] = pack(Qgu).cpu(); tensors[p + "gate_up_scales"] = Sgu.float().cpu()
    tensors[p + "down_qweight"] = pack(Qdn).cpu(); tensors[p + "down_scales"] = Sdn.float().cpu()
    assert torch.isfinite(Sgu).all() and (Sgu > 0).all() and torch.isfinite(Sdn).all() and (Sdn > 0).all()
    tensors = {k: v.cpu().contiguous() for k, v in tensors.items()}
    save_file(tensors, f"{OUT}/{f}.tmp", metadata={"format": "pt"})
    os.replace(f"{OUT}/{f}.tmp", f"{OUT}/{f}")
    stats[li] = dict(gate_up_gptq=e_gu, gate_up_rtn=e_gu_rtn, down_gptq=e_dn, down_rtn=e_dn_rtn,
                     min_tokens=int(cnt.min()), rtn_experts=int((cnt < 512).sum()), median_tokens=int(cnt.median()), seconds=round(time.time() - t, 1))
    log(f"layer {li}: {stats[li]}")
    os.remove(f"{SRC}/{f}")
    return list(tensors)


def main():
    stats_path = f"{WORK}/quant_stats.json"
    stats = json.load(open(stats_path)) if os.path.exists(stats_path) else {}
    stats = {int(k): v for k, v in stats.items()}
    calib = reference_and_calibrate()
    failed = []
    cfg = json.load(open(f"{SRC}/config.json"))
    for li in range(cfg["first_k_dense_replace"], cfg["num_hidden_layers"]):
        if os.path.exists(f"{OUT}/model-layer-{li:02d}.safetensors"):
            continue
        X, ids = calib.pop(li)
        try:
            quantize_layer(li, X, ids, stats)
        except Exception:
            import traceback; traceback.print_exc()
            log(f"layer {li} FAILED; continuing (rerun to retry)"); failed.append(li)
        json.dump(stats, open(stats_path, "w"), indent=1)
        del X, ids; gc.collect(); torch.cuda.empty_cache()
    if failed:
        log(f"FAILED layers: {failed}"); sys.exit(1)
    # everything else is copied unchanged
    for f in os.listdir(SRC):
        if f.startswith(".") or f in ("prism.py", "example.py", "requirements.txt", "README.md", "model.safetensors.index.json"):
            continue
        if not os.path.exists(f"{OUT}/{f}"):
            (shutil.copytree if os.path.isdir(f"{SRC}/{f}") else shutil.copy2)(f"{SRC}/{f}", f"{OUT}/{f}")
    wm = {}
    for f in sorted(os.listdir(OUT)):
        if f.endswith(".safetensors") and f.startswith("model-"):
            with safe_open(f"{OUT}/{f}", "pt") as sf:
                for k in sf.keys(): wm[k] = f
    json.dump({"metadata": {}, "weight_map": wm}, open(f"{OUT}/model.safetensors.index.json", "w"), indent=1)
    cfg["quantization"] = {
        "routed_experts": {"format": "int4", "method": "GPTQ (damp 0.01, MSE clip search)", "symmetric": True,
                           "group_size": GROUP, "scale_dtype": "float32",
                           "packing": "uint8 [E, N, K/2]; in each 128-group byte j = elem j (low nibble) | elem j+64 (high nibble); value = nibble - 8",
                           "calibration": f"{CALIB_TOKENS} tokens of Prism-format requests over multilingual Wikipedia (11 languages)"},
        "other_projections": {"format": "float8_e4m3fn", "weight_scale": "per output row (absmax / 448)", "note": "unchanged from sarvamai/prism"},
        "embeddings_norms": "bfloat16", "router": "float32"}
    json.dump(cfg, open(f"{OUT}/config.json", "w"), indent=2)
    log("done")


if __name__ == "__main__":
    main()