"""Quantize the vocabulary embedding of the laya ONNX graph, keeping the int4 layer weights intact.

Motivation: MatMulNBits only rewrites MatMul weights, so `tok_embeddings` [256000, 768] stays fp32
= 786MB = 61% of the model. ORT's own dynamic quantization does cover Gather, but with a SINGLE
per-tensor scale for the whole 256k x 768 table (that is part of why the int8 arm lost 8pp).

This does graph surgery instead: replace the fp32 embedding initializer with a quantized one plus a
DequantizeLinear feeding the Gather, with a scale PER ROW (每行一个 scale = 每个 token 向量一个).

  python emb_quant.py --src laya-zh-int4-b32.onnx --dst laya-zh-int4-b32-emb8r.onnx --mode int8-row
  modes: int8-row (per-row int8) | fp8-row (E4M3 per-row, needs opset 21 + ORT fp8 kernel)
         u4-row (per-row uint4)   | u4-block32 (4-bit blocked along K, ONNX blocked-quant layout)

The embedding is a pure lookup table, so quantization error here = a small perturbation of the
token vector; per-row scales keep that perturbation proportional to the row's own magnitude.
"""
import argparse, os, time

import numpy as np
import onnx
from onnx import TensorProto, helper, numpy_helper

EMB = "encoder.embeddings.tok_embeddings.weight"
DQ_OUT = EMB + "_dqout"   # the DequantizeLinear result the Gather must consume


def get_emb(model):
    for t in model.graph.initializer:
        if t.name == EMB:
            return numpy_helper.to_array(t)
    raise SystemExit(f"{EMB} not found")


def replace_emb(model, new_inits):
    """Insert DequantizeLinear(<q>, <scale>, [zp]) -> replaces the fp32 initializer as Gather input."""
    # drop the old initializer
    old = [i for i in model.graph.initializer if i.name == EMB]
    for i in old:
        model.graph.initializer.remove(i)
    model.graph.initializer.extend(new_inits)
    # find the Gather that consumes EMB
    for n in model.graph.node:
        if n.op_type == "Gather" and EMB in n.input:
            idx = list(n.input).index(EMB)
            n.input[idx] = DQ_OUT
            return n
    raise SystemExit("no Gather consumes the embedding")


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--src", required=True)
    ap.add_argument("--dst", required=True)
    ap.add_argument("--mode", default="int8-row")
    ap.add_argument("--block", type=int, default=32)
    ap.add_argument("--axis", type=int, default=1,
                    help="1 = per-column scale [D] (ORT default per-axis), 0 = per-row scale [V]")
    a = ap.parse_args()
    src, dst = os.path.expanduser(a.src), os.path.expanduser(a.dst)
    t0 = time.perf_counter()
    model = onnx.load(src, load_external_data=True)
    if a.mode.startswith("u4"):
        del model.graph.value_info[:]   # packed/blocked layouts need ORT to re-infer shapes
    w = get_emb(model).astype(np.float32)
    N, K = w.shape
    print(f"embedding {w.shape} {w.nbytes/1e6:.0f}MB mode={a.mode}", flush=True)

    new_inits, dq_inputs, attrs, extra = [], [], {}, {}
    if a.mode == "int8-row":
        if a.axis == 1:                      # one scale per hidden dim (ORT's per-axis default)
            sc = (np.abs(w).max(axis=0) / 127.0).astype(np.float32)                 # [D]
            q = np.clip(np.round(w / sc[None, :]), -127, 127).astype(np.int8)
        else:                                # one scale per token row
            sc = (np.abs(w).max(axis=1) / 127.0).astype(np.float32)                 # [V]
            q = np.clip(np.round(w / sc[:, None]), -127, 127).astype(np.int8)
            attrs["axis"] = 0
        dq_inputs = [EMB + "_q8", EMB + "_scale"]
        new_inits = [numpy_helper.from_array(q, EMB + "_q8"),
                     numpy_helper.from_array(sc, EMB + "_scale")]
    elif a.mode == "fp8-row":
        import ml_dtypes
        if a.axis == 1:
            sc = (np.abs(w).max(axis=0) / 448.0).astype(np.float32)                 # E4M3 max 448
            q = (w / sc[None, :]).astype(ml_dtypes.float8_e4m3fn)
        else:
            sc = (np.abs(w).max(axis=1) / 448.0).astype(np.float32)
            q = (w / sc[:, None]).astype(ml_dtypes.float8_e4m3fn)
            attrs["axis"] = 0
        dq_inputs = [EMB + "_f8", EMB + "_scale"]
        new_inits = [numpy_helper.from_array(q, EMB + "_f8"),
                     numpy_helper.from_array(sc, EMB + "_scale")]
        extra["opset"] = 21                            # float8 types land in 19+, DQ/blocked in 21
    elif a.mode == "u4-row":
        sc = np.abs(w).max(axis=1) / 7.0
        q = np.clip(np.round(w / sc[:, None]) + 8, 0, 15).astype(np.uint8)   # zp=8 -> symmetric
        packed = (q[:, 0::2] | (q[:, 1::2] << 4)).astype(np.uint8)           # low nibble = first elem
        zp = np.full((N, 1), 8, np.uint8)
        dq_inputs = [EMB + "_q4", EMB + "_scale", EMB + "_zp"]
        new_inits = [numpy_helper.from_array(packed, EMB + "_q4"),
                     numpy_helper.from_array(sc.reshape(N, 1).astype(np.float32), EMB + "_scale"),
                     numpy_helper.from_array(zp.reshape(N, 1), EMB + "_zp")]
        extra["opset"] = 21
        attrs["block_size"] = K          # one block == the whole row
    elif a.mode == "u4-gbq":
        # 4-bit vocabulary via com.microsoft::GatherBlockQuantized.
        # Layout (verified on this box with a synthetic probe): data uint8 [V, D/2] with the LOW
        # nibble = first element; scales fp32 [V, D/block]; zero_points are OPTIONAL for 4-bit and
        # when omitted the op uses 8 (symmetric) -- exactly what we want, and it keeps the table
        # packed at inference (no fp32 materialisation, unlike DequantizeLinear).
        B = a.block
        assert B >= 16 and (B & (B - 1)) == 0 and K % B == 0, "block_size must be a power of 2 >= 16 dividing D"
        wb = w.reshape(N, K // B, B)
        sc = np.abs(wb).max(axis=2) / 7.0
        q = np.clip(np.round(wb / sc[:, :, None]) + 8, 0, 15).astype(np.uint8)
        packed = (q[:, :, 0::2] | (q[:, :, 1::2] << 4)).reshape(N, K // 2)
        new_inits = [numpy_helper.from_array(packed.astype(np.uint8), EMB + "_q4packed"),
                     numpy_helper.from_array(sc.astype(np.float32), EMB + "_scales")]
        gbq_attrs = dict(bits=4, block_size=B, gather_axis=0, quantize_axis=1)
    elif a.mode == "u4-block32":
        B = a.block
        assert K % B == 0 and B % 2 == 0
        wb = w.reshape(N, K // B, B)
        sc = np.abs(wb).max(axis=2) / 7.0
        q = np.clip(np.round(wb / sc[:, :, None]) + 8, 0, 15).astype(np.uint8)
        packed = (q[:, :, 0::2] | (q[:, :, 1::2] << 4)).reshape(N, K // 2)
        zp = np.full((N, K // B), 8, np.uint8)
        dq_inputs = [EMB + "_q4", EMB + "_scale", EMB + "_zp"]
        new_inits = [numpy_helper.from_array(packed, EMB + "_q4"),
                     numpy_helper.from_array(sc.astype(np.float32), EMB + "_scale"),
                     numpy_helper.from_array(zp, EMB + "_zp")]
        extra["opset"] = 21
        attrs["block_size"] = B
    else:
        raise SystemExit("unknown mode")

    if a.mode == "u4-gbq":
        gathered = [n for n in model.graph.node if n.op_type == "Gather" and EMB in n.input][0]
        idx_in = list(gathered.input)[1 - list(gathered.input).index(EMB)]
        node = helper.make_node("GatherBlockQuantized", [EMB + "_q4packed", idx_in, EMB + "_scales"],
                                list(gathered.output), name="emb_gbq", domain="com.microsoft", **gbq_attrs)
        model.graph.node.remove(gathered)
        model.graph.node.append(node)
        for i in [t for t in model.graph.initializer if t.name == EMB]:
            model.graph.initializer.remove(i)
        model.graph.initializer.extend(new_inits)
        doms = {op.domain for op in model.opset_import}
        if "com.microsoft" not in doms:
            model.opset_import.append(helper.make_opsetid("com.microsoft", 1))
        g = node
    else:
        dq = helper.make_node("DequantizeLinear", dq_inputs, [DQ_OUT], name="emb_dq", **attrs)
        g = replace_emb(model, new_inits)
        model.graph.node.insert(list(model.graph.node).index(g), dq)
    if extra.get("opset"):
        for op in model.opset_import:
            if op.domain in ("", "ai.onnx"):
                op.version = max(op.version, extra["opset"])
    for stale in (dst, dst + ".data"):      # save_model refuses to overwrite an existing .data
        if os.path.exists(stale):
            os.remove(stale)
    onnx.save_model(model, dst, save_as_external_data=True, all_tensors_to_one_file=True,
                    location=os.path.basename(dst) + ".data", size_threshold=1024)
    got = os.path.getsize(dst) + (os.path.getsize(dst + ".data") if os.path.exists(dst + ".data") else 0)
    print(f"done in {time.perf_counter()-t0:.1f}s -> {got/1e6:.0f}MB  (embedding part: "
          f"{sum(len(i.raw_data) for i in new_inits)/1e6:.0f}MB, was {w.nbytes/1e6:.0f}MB)")
    print("EMB_QUANT_DONE")


if __name__ == "__main__":
    main()