"""int4 (weight-only, block-wise) quantization of the laya ONNX graph via ONNX Runtime's
MatMulNBits path — the same mechanism ORT/GenAI use for 4-bit LLM weights.

  python int4_quantize.py --src laya-zh-fp32-clean.onnx --dst laya-zh-int4-b32.onnx \
      --block-size 32 --symmetric

Only `MatMul` nodes whose weight input is a constant get replaced by `com.microsoft::MatMulNBits`
(attention QK^T / softmax·V keep their fp32 kernel). No calibration data needed (RTN).
"""
import argparse, collections, os, time

import onnx
from onnxruntime.quantization.matmul_nbits_quantizer import (
    DefaultWeightOnlyQuantConfig, MatMulNBitsQuantizer)


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--src", required=True)
    ap.add_argument("--dst", required=True)
    ap.add_argument("--bits", type=int, default=4)
    ap.add_argument("--block-size", type=int, default=32)
    ap.add_argument("--symmetric", action="store_true")
    ap.add_argument("--accuracy-level", type=int, default=None)
    a = ap.parse_args()
    src, dst = os.path.expanduser(a.src), os.path.expanduser(a.dst)

    cfg = DefaultWeightOnlyQuantConfig(block_size=a.block_size, is_symmetric=a.symmetric,
                                       bits=a.bits, accuracy_level=a.accuracy_level)
    print(f"int{a.bits} weight-only: block_size={a.block_size} symmetric={a.symmetric} "
          f"accuracy_level={a.accuracy_level} | {src} -> {dst}", flush=True)
    t0 = time.perf_counter()
    q = MatMulNBitsQuantizer(model=src, algo_config=cfg)   # first positional arg is `model`
    print("config loaded, quantizing ...", flush=True)
    q.process()
    q.model.save_model_to_file(dst, use_external_data_format=True)
    dt = time.perf_counter() - t0

    m = onnx.load(dst, load_external_data=False)
    ops = collections.Counter(n.op_type for n in m.graph.node)
    size = os.path.getsize(dst) + (os.path.getsize(dst + ".data") if os.path.exists(dst + ".data") else 0)
    src_size = os.path.getsize(src) + (os.path.getsize(src + ".data") if os.path.exists(src + ".data") else 0)
    print(f"done in {dt:.1f}s | {src_size/1e6:.0f}MB -> {size/1e6:.0f}MB ({src_size/size:.1f}x smaller)")
    print(f"node mix: MatMul={ops.get('MatMul',0)} MatMulNBits={ops.get('MatMulNBits',0)} "
          f"Gemm={ops.get('Gemm',0)} total={sum(ops.values())}")
    print("INT4_QUANTIZE_DONE")


if __name__ == "__main__":
    main()