#!/usr/bin/env python3
"""Calibrated weight-only quantization of NadirNet into ONNX MatMulNBits (8-bit and 4-bit).

Layer-sequential GPTQ: each layer's Hessian comes from real crawled pages flowing through the already
quantized earlier layers. Each group of 32 weights gets asymmetric uint codes whose clipping range is
chosen by an MSE search. Codes are packed in ORT's MatMulNBits layout, and every variant is checked
three ways: against a PyTorch fake-quant copy (packing parity), against fp32 (fidelity) and against
the Haiku targets (quality)."""
import argparse, copy, json, os, pickle, time
import numpy as np
import torch
import torch.nn as nn

from model import NadirNet
from train import prep
from export import feeds_for, elem_feeds, fidelity, bench, NAMES_IN

GS = 32


class Stop(Exception):
    pass


def group_params(w, bits, maxshrink, grid=100):
    """Per-row scale/zero-point for one group (rows, GS), minimizing |q(w)-w|^2.4 over clipping ranges."""
    maxq = 2 ** bits - 1
    xmin = w.min(1).values.clamp(max=0)
    xmax = w.max(1).values.clamp(min=0)
    best = torch.full_like(xmin, float("inf"))
    bs, bz = torch.ones_like(xmin), torch.zeros_like(xmin)
    for i in range(int(maxshrink * grid) + 1):
        p = 1 - i / grid
        lo, hi = p * xmin, p * xmax
        s = ((hi - lo) / maxq).clamp(min=1e-10)
        z = torch.round(-lo / s).clamp(0, maxq)
        q = (torch.clamp(torch.round(w / s[:, None]) + z[:, None], 0, maxq) - z[:, None]) * s[:, None]
        err = (q - w).abs().pow(2.4).sum(1)
        m = err < best
        best[m], bs[m], bz[m] = err[m], s[m], z[m]
    return bs, bz


def gptq(W, H, bits, percdamp=0.01, blocksize=128, gs=GS):
    """W (out, Kpad) with Kpad % gs == 0, H (Kpad, Kpad); gs=None quantizes per output channel (one group per row).
    Returns dequantized W, codes, scales, zero points."""
    maxq = 2 ** bits - 1
    maxshrink = 0.5 if bits == 4 else 0.2
    W = W.clone().double()
    H = H.clone().double()
    rows, cols = W.shape
    dead = torch.diag(H) == 0
    H[dead, dead] = 1
    W[:, dead] = 0
    H += percdamp * torch.mean(torch.diag(H)) * torch.eye(cols, dtype=H.dtype)
    Hinv = torch.linalg.cholesky(torch.cholesky_inverse(torch.linalg.cholesky(H)), upper=True)
    Q = torch.zeros_like(W)
    codes = torch.zeros(rows, cols, dtype=torch.uint8)
    if gs is None:  # per channel: one range per output row, chosen on the whole row up front
        s, z = group_params(W, bits, maxshrink)
        S, Z = s[:, None].clone(), z[:, None].clone()
    else:
        S = torch.zeros(rows, cols // gs, dtype=torch.float64)
        Z = torch.zeros(rows, cols // gs, dtype=torch.float64)
    for i1 in range(0, cols, blocksize):
        i2 = min(i1 + blocksize, cols)
        W1 = W[:, i1:i2].clone()
        E1 = torch.zeros_like(W1)
        H1 = Hinv[i1:i2, i1:i2]
        for i in range(i2 - i1):
            c = i1 + i
            if gs is not None and c % gs == 0:  # groups never straddle blocks (blocksize % gs == 0)
                s, z = group_params(W1[:, i:i + gs], bits, maxshrink)
                S[:, c // gs], Z[:, c // gs] = s, z
            w = W1[:, i]
            qc = torch.clamp(torch.round(w / s) + z, 0, maxq)
            q = (qc - z) * s
            codes[:, c] = qc.to(torch.uint8)
            Q[:, c] = q
            err = (w - q) / H1[i, i]
            W1[:, i:] -= err[:, None] * H1[i, i:][None, :]
            E1[:, i] = err
        W[:, i2:] -= E1 @ Hinv[i1:i2, i2:]
    return Q.float(), codes, S.float(), Z.to(torch.uint8)


def pack(codes, S, Z, bits):
    """MatMulNBits layout: B (N, k_blocks, blob), scales (N, k_blocks), zero points (N, ceil(k_blocks*bits/8))."""
    c = codes.numpy()
    N, Kp = c.shape
    kb = Kp // GS
    z = Z.numpy().astype(np.uint8)
    if bits == 8:
        B = c.reshape(N, kb, GS).astype(np.uint8)
        zp = z
    else:
        c = c.reshape(N, kb, GS)
        B = (c[..., 0::2] | (c[..., 1::2] << 4)).astype(np.uint8)
        if kb % 2:
            z = np.concatenate([z, np.full((N, 1), 8, np.uint8)], 1)
        zp = (z[:, 0::2] | (z[:, 1::2] << 4)).astype(np.uint8)
    return B, S.numpy().astype(np.float32), zp


def linear_order(model, sample):
    """Linear modules in the order the forward pass calls them."""
    order, hooks = [], []
    for name, mod in model.named_modules():
        if isinstance(mod, nn.Linear):
            hooks.append(mod.register_forward_pre_hook(lambda m, a, n=name: order.append(n) if n not in order else None))
    with torch.no_grad():
        model(*[torch.from_numpy(sample[k]) for k in NAMES_IN])
    for h in hooks:
        h.remove()
    return order


def onnx_matmuls(model, onnx_model):
    """Map each Linear to the MatMul node whose initializer is its transposed weight (Gemm exports stay fp32)."""
    from onnx import numpy_helper
    inits = {i.name: i for i in onnx_model.graph.initializer}
    by_shape = {}
    for node in onnx_model.graph.node:
        if node.op_type == "MatMul" and node.input[1] in inits:
            arr = numpy_helper.to_array(inits[node.input[1]])
            by_shape.setdefault(arr.shape, []).append((node, arr))
    out = {}
    for name, mod in model.named_modules():
        if isinstance(mod, nn.Linear):
            wt = mod.weight.detach().numpy().T
            for node, arr in by_shape.get(wt.shape, []):
                if np.array_equal(arr, wt):
                    out[name] = node
                    break
    return out


def hessian(model, name, calib):
    """2/n * X^T X of the inputs to one Linear, over all real token rows of the calibration pages."""
    mod = dict(model.named_modules())[name]
    K = mod.in_features
    H = torch.zeros(K, K, dtype=torch.float64)
    n = [0]

    def hook(m, a):
        x = a[0].detach().reshape(-1, K).double()
        H.add_(x.T @ x)
        n[0] += x.shape[0]
        raise Stop

    h = mod.register_forward_pre_hook(hook)
    with torch.no_grad():
        for f in calib:
            try:
                model(*[torch.from_numpy(f[k]) for k in NAMES_IN])
            except Stop:
                pass
    h.remove()
    return H * (2.0 / max(1, n[0]))


def quantize_model(model, calib, bits_for, order):
    """Sequential GPTQ over `order`; returns the fake-quant model and packed tensors per layer."""
    qm = copy.deepcopy(model).eval()
    mods = dict(qm.named_modules())
    packed = {}
    for name in order:
        bits = bits_for.get(name)
        if not bits:
            continue
        mod = mods[name]
        K = mod.in_features
        H = hessian(qm, name, calib)
        t = time.time()
        # blocks of 32 along K (zero-padded to a multiple of 32)
        Kp = -(-K // GS) * GS
        Hp = torch.zeros(Kp, Kp, dtype=torch.float64); Hp[:K, :K] = H
        Wp = torch.zeros(mod.out_features, Kp); Wp[:, :K] = mod.weight.detach()
        Q, codes, S, Z = gptq(Wp, Hp, bits)
        if bits == 8:
            # 8-bit: uint8 codes as a (K/32, 32, N) tensor with per-block scale and zero point, dequantized in-graph
            N, kb = mod.out_features, Kp // GS
            packed[name] = ("dq8", K, codes.numpy().T.reshape(kb, GS, N).copy(), S.numpy().T[:, None, :].astype(np.float32).copy(),
                            Z.numpy().T[:, None, :].astype(np.uint8).copy())
        else:
            # 4-bit: packed as MatMulNBits
            packed[name] = ("nbits", K) + pack(codes, S, Z, bits)
        with torch.no_grad():
            mod.weight.copy_(Q[:, :K])
        print(f"  {name} {tuple(mod.weight.shape)} {bits}-bit  {time.time() - t:.1f}s", flush=True)
    return qm, packed


def build_onnx(fp32_path, mapping, packed, out_path, accuracy_level=0):
    import onnx
    from onnx import helper, numpy_helper
    m = onnx.load(fp32_path)
    g = m.graph
    drop_inits = set()
    for name, (kind, K, B, S, zp) in packed.items():
        node = mapping[name]
        wname = node.input[1]
        idx = list(g.node).index(node)
        if kind == "dq8":
            # uint8 codes dequantized with plain Cast/Sub/Mul: every onnxruntime constant-folds this into exact fp32
            # weights at load. (DequantizeLinear->MatMul gets rewritten by onnxruntime into an activation-quantized
            # kernel, which costs accuracy.)
            kb, gs, N = B.shape
            inits = [numpy_helper.from_array(B, wname + "_q8"), numpy_helper.from_array(zp, wname + "_zero_point"),
                     numpy_helper.from_array(S, wname + "_scale"), numpy_helper.from_array(np.array([kb * gs, N], np.int64), wname + "_kn")]
            g.initializer.extend(inits)
            base = (node.name or wname) + "_deq8"
            F = onnx.TensorProto.FLOAT
            deq = [helper.make_node("Cast", [inits[0].name], [wname + "_f"], name=base + "_cast", to=F),
                   helper.make_node("Cast", [inits[1].name], [wname + "_zf"], name=base + "_castz", to=F),
                   helper.make_node("Sub", [wname + "_f", wname + "_zf"], [wname + "_c"], name=base + "_sub"),
                   helper.make_node("Mul", [wname + "_c", inits[2].name], [wname + "_d3"], name=base + "_mul"),
                   helper.make_node("Reshape", [wname + "_d3", inits[3].name], [wname + ("_dq" if kb * gs == K else "_dp")], name=base + "_reshape")]
            if kb * gs != K:  # drop the zero padding rows
                lims = [numpy_helper.from_array(np.array(v, np.int64), wname + n) for v, n in (([0], "_s0"), ([K], "_sK"), ([0], "_ax"))]
                g.initializer.extend(lims)
                deq.append(helper.make_node("Slice", [wname + "_dp"] + [i.name for i in lims], [wname + "_dq"], name=base + "_slice"))
            g.node[idx].input[1] = wname + "_dq"  # `node` belongs to another copy of the graph: edit this one
            for k, dn in enumerate(deq):
                g.node.insert(idx + k, dn)
        else:
            N = B.shape[0]
            inits = [numpy_helper.from_array(B, wname + "_Q4"), numpy_helper.from_array(S, wname + "_scales"),
                     numpy_helper.from_array(zp, wname + "_zero_points")]
            g.initializer.extend(inits)
            kw = dict(K=K, N=N, bits=4, block_size=GS)
            if accuracy_level:
                kw["accuracy_level"] = accuracy_level
            new = helper.make_node("MatMulNBits", [node.input[0]] + [i.name for i in inits], list(node.output),
                                   name=(node.name or wname) + "_Q4", domain="com.microsoft", **kw)
            g.node.remove(node)
            g.node.insert(idx, new)
        drop_inits.add(wname)
    keep = [i for i in g.initializer if i.name not in drop_inits]
    del g.initializer[:]
    g.initializer.extend(keep)
    if any(n.domain == "com.microsoft" for n in g.node) and not any(o.domain == "com.microsoft" for o in m.opset_import):
        m.opset_import.append(helper.make_opsetid("com.microsoft", 1))
    onnx.checker.check_model(m)
    onnx.save(m, out_path)


def quality(sess, samples, esess=None, ref=None):
    """dE vs the training targets (Haiku for src 0, native dark for src 1); with the element graph, how often the
    element decisions computed from this page graph's outputs match those from the fp32 page graph."""
    agg = {0: [0.0, 0.0], 1: [0.0, 0.0]}
    ok = n = 0
    for s in samples:
        f = feeds_for(s)
        lab, _, ctx = sess.run(None, f)
        k = len(s["role"])
        m = s["tmask"][:k]
        if m.any():
            d = np.linalg.norm(lab[0, :k] - s["lab_t"][:k], axis=-1)
            w = s["w"][:k] * m
            agg[int(s["src"])][0] += float((d * w).sum() / max(w.sum(), 1e-9)); agg[int(s["src"])][1] += 1
        if esess is not None and ref is not None and len(s["E_ut"]):
            rl, _, rc = ref.run(None, f)
            a = esess.run(None, elem_feeds(s, lab, ctx)[0])[0][0].argmax(-1)
            b = esess.run(None, elem_feeds(s, rl, rc)[0])[0][0].argmax(-1)
            ok += int((a == b).sum()); n += len(a)
    out = {"haiku_dE": agg[0][0] / max(1, agg[0][1]), "native_dE": agg[1][0] / max(1, agg[1][1])}
    if n:
        out["elem_decisions_same_as_fp32"] = ok / n
    return out


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--ckpt", required=True)
    ap.add_argument("--data", required=True)
    ap.add_argument("--fp32", required=True, help="fp32 ONNX from export.py")
    ap.add_argument("--out", required=True)
    ap.add_argument("--calib", type=int, default=256)
    ap.add_argument("--eval", type=int, default=400)
    ap.add_argument("--threads", type=int, default=8)
    ap.add_argument("--elem", default="", help="element graph (fp32) to measure decision agreement")
    args = ap.parse_args()
    torch.set_num_threads(args.threads)
    os.makedirs(args.out, exist_ok=True)
    import onnx, onnxruntime as ort
    ck = torch.load(args.ckpt, map_location="cpu")
    model = NadirNet(**ck["cfg"]); model.load_state_dict(ck["model"]); model.eval()
    data = pickle.load(open(args.data, "rb"))
    hv = [s for s in data["val"] if s["src"] == 0]; nv = [s for s in data["val"] if s["src"] == 1]
    val = [prep(s) for s in hv[: args.eval * 3 // 4] + nv[: args.eval - min(len(hv), args.eval * 3 // 4)]]  # both sources
    rng = np.random.default_rng(0)
    pick = rng.choice(len(data["train"]), size=min(args.calib, len(data["train"])), replace=False)
    calib = [feeds_for(prep(data["train"][i])) for i in pick]
    order = linear_order(model, calib[0])
    mapping = onnx_matmuls(model, onnx.load(args.fp32))
    quantizable = [n for n in order if n in mapping and not n.endswith(".rel")]
    print(f"{len(quantizable)}/{len(order)} linears map to MatMul; fp32 kept: {[n for n in order if n not in quantizable]}", flush=True)
    edges = [n for n in quantizable if n.startswith(("tok_in.", "img_in.", "color_head.", "ctx_head."))]
    plans = {"int8": {n: 8 for n in quantizable},
             "int4": {n: (8 if n in edges else 4) for n in quantizable}}
    so = ort.SessionOptions(); so.intra_op_num_threads = 4
    ref = ort.InferenceSession(args.fp32, so, providers=["CPUExecutionProvider"])
    report = {"format": {"8-bit layers": f"uint8 blocks of {GS} with fp32 scale + uint8 zero point, dequantized in-graph (Cast/Sub/Mul, folded at load)", "4-bit layers": f"uint4 blocks of {GS}, MatMulNBits"},
              "group_size": GS, "calib_pages": len(calib), "eval_pages": len(val), "fp32_layers": [n for n in order if n not in quantizable],
              "int4_8bit_edge_layers": edges, "variants": {}}
    esess = ort.InferenceSession(args.elem, so, providers=["CPUExecutionProvider"]) if args.elem else None
    bfeeds = [feeds_for(s) for s in val[:20]]
    r = quality(ref, val); r["size_mb"] = round(os.path.getsize(args.fp32) / 1e6, 2); r["cpu_ms_median"] = round(bench(ref, bfeeds), 2)
    report["variants"]["fp32"] = r
    print("fp32", r, flush=True)
    for vname, bits_for in plans.items():
        print(f"== {vname}", flush=True)
        qm, packed = quantize_model(model, calib, bits_for, quantizable)
        path = os.path.join(args.out, f"nadir-{vname}.onnx")
        build_onnx(args.fp32, mapping, packed, path)
        sess = ort.InferenceSession(path, so, providers=["CPUExecutionProvider"])
        # packing parity: ORT MatMulNBits graph vs the PyTorch fake-quant model
        par = 0.0
        with torch.no_grad():
            for s in val[:8]:
                f = feeds_for(s)
                tq = qm(*[torch.from_numpy(f[k]) for k in NAMES_IN])[0].numpy()
                par = max(par, float(np.abs(tq - sess.run(None, f)[0]).max()))
        r = fidelity(sess, ref, val)
        r.update(quality(sess, val, esess, ref))
        r["pack_parity_max_abs"] = par
        r["size_mb"] = round(os.path.getsize(path) / 1e6, 2)
        r["cpu_ms_median"] = round(bench(sess, bfeeds), 2)
        report["variants"][vname] = r
        print(vname, r, flush=True)
        torch.save({"cfg": ck["cfg"], "model": qm.state_dict()}, os.path.join(args.out, f"fakequant-{vname}.pt"))
    json.dump(report, open(os.path.join(args.out, "quant_report.json"), "w"), indent=1)


if __name__ == "__main__":
    main()