"""Quantize SAM3 ONNX components to int4 via MatMul4BitsQuantizer.

This is the most aggressive quantization step. Targets specifically the MatMul
(matrix multiply) operations — which dominate transformer compute — and stores
their weights as 4-bit integers grouped into blocks. Other ops (LayerNorm,
activations, FPN convolutions) stay at fp32.

Block-wise int4 is the same scheme used by:
  - HuggingFace's `dtype: "q4"` in transformers.js
  - GPTQ and AWQ for LLMs
  - onnx-community's SAM3-Tracker quantization

Expected: ~8x smaller than fp32, small but real accuracy hit. Validate before
using in production.

Run validate_sam3_e2e.py against the int4 files after to confirm.
"""

from pathlib import Path
import time

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

OUTPUT_DIR = Path("sam3-onnx-test")

COMPONENTS = [
    "vision_encoder.onnx",
    "text_encoder.onnx",
    "decoder.onnx",
]


def quantize_one(name: str) -> None:
    in_path = OUTPUT_DIR / name
    out_name = name.replace(".onnx", "_int4.onnx")
    out_path = OUTPUT_DIR / out_name

    print(f"\n[{name}] → [{out_name}]")
    if not in_path.exists():
        print(f"  ❌ {in_path} not found")
        return

    in_size = in_path.stat().st_size
    in_data = OUTPUT_DIR / f"{name}.data"
    if in_data.exists():
        in_size += in_data.stat().st_size
    print(f"  Source size: {in_size / 1024 / 1024:.1f} MB total")

    print(f"  Loading model ...")
    t0 = time.time()
    model = onnx.load(str(in_path), load_external_data=True)

    print(f"  Running MatMul int4 quantization (block_size=128, symmetric) ...")
    # Newer onnxruntime uses a config object instead of constructor kwargs.
    # DefaultWeightOnlyQuantConfig is the Round-To-Nearest scheme — simple and reliable.
    config = DefaultWeightOnlyQuantConfig(
        block_size=128,
        is_symmetric=True,
        accuracy_level=4,      # int8 kernel at inference time (fastest)
        bits=4,                # the "N" in "N-bits"
    )
    quantizer = MatMulNBitsQuantizer(model, algo_config=config)
    quantizer.process()

    print(f"  Saving to {out_path} ...")
    # Clean any existing output files
    for p in [out_path, OUTPUT_DIR / f"{out_name}.data"]:
        if p.exists():
            p.unlink()

    # Save with external data for the big ones (vision/text encoders may still
    # exceed 2GB after quantization in edge cases — easier to always allow it)
    quantized_model_proto = quantizer.model.model
    is_big = in_size > 200 * 1024 * 1024
    if is_big:
        onnx.save(
            quantized_model_proto,
            str(out_path),
            save_as_external_data=True,
            all_tensors_to_one_file=True,
            location=f"{out_name}.data",
            size_threshold=1024,
        )
    else:
        onnx.save(quantized_model_proto, str(out_path))

    elapsed = time.time() - t0

    out_size = out_path.stat().st_size
    out_data = OUTPUT_DIR / f"{out_name}.data"
    if out_data.exists():
        out_size += out_data.stat().st_size

    print(f"  Output size: {out_size / 1024 / 1024:.1f} MB total")
    print(f"  Reduction:   {in_size / out_size:.2f}x  ({elapsed:.1f}s)")


def main() -> None:
    print("SAM3 fp32 → int4 (MatMul block-quantize)\n" + "=" * 50)
    for name in COMPONENTS:
        try:
            quantize_one(name)
        except Exception as e:
            print(f"  ❌ FAILED: {type(e).__name__}: {e}")
            import traceback
            traceback.print_exc()

    print("\n" + "=" * 50)
    print("Done. New int4 files:")
    for f in sorted(OUTPUT_DIR.iterdir()):
        if "int4" in f.name:
            size = f.stat().st_size / 1024 / 1024
            print(f"  {f.name}: {size:.1f} MB")


if __name__ == "__main__":
    main()