"""
Quantize exported Cohere Transcribe ONNX models to INT8 and optionally INT4.

Requirements:
    pip install onnx onnxruntime

Usage:
    python quantize.py --input-dir ./onnx-export
    python quantize.py --input-dir ./onnx-export --int4   # also produce INT4 variants

Output:
    cohere-encoder.int8.onnx
    cohere-decoder.int8.onnx
    cohere-encoder.int4.onnx  (if --int4)
    cohere-decoder.int4.onnx  (if --int4)
"""

import argparse
import os
from pathlib import Path


def quantize_int8(input_path, output_path, nodes_to_exclude=None, use_external_data=False):
    """Dynamic INT8 quantization -- best for transformer/attention models."""
    from onnxruntime.quantization import quantize_dynamic, QuantType

    print(f"  INT8: {input_path} -> {output_path}")
    quantize_dynamic(
        model_input=input_path,
        model_output=output_path,
        weight_type=QuantType.QInt8,
        nodes_to_exclude=nodes_to_exclude or [],
        use_external_data_format=use_external_data,
    )
    if os.path.exists(output_path):
        quantized_size = os.path.getsize(output_path) / (1024 * 1024)
        print(f"  Output: {quantized_size:.1f} MB (graph only, weights external)" if use_external_data else f"  Output: {quantized_size:.1f} MB")


def quantize_int4(input_path, output_path):
    """INT4 weight-only quantization -- aggressive, ~8x compression."""
    from onnxruntime.quantization import matmul_4bits_quantizer, quant_utils

    print(f"  INT4: {input_path} -> {output_path}")
    quant_config = matmul_4bits_quantizer.DefaultWeightOnlyQuantConfig(
        block_size=128,
        is_symmetric=True,
        accuracy_level=4,
        quant_format=quant_utils.QuantFormat.QOperator,
        op_types_to_quantize=("MatMul",),
    )

    model = quant_utils.load_model_with_shape_infer(Path(input_path))
    quant = matmul_4bits_quantizer.MatMul4BitsQuantizer(
        model,
        nodes_to_exclude=None,
        nodes_to_include=None,
        algo_config=quant_config,
    )
    quant.process()
    # Use external data format for large models (encoder is >2GB)
    quant.model.save_model_to_file(output_path, True)

    if os.path.exists(output_path):
        quantized_size = os.path.getsize(output_path) / (1024 * 1024)
        print(f"  Output: {quantized_size:.1f} MB")


def main():
    parser = argparse.ArgumentParser(description="Quantize Cohere Transcribe ONNX models")
    parser.add_argument("--input-dir", type=str, default="./onnx-export", help="Directory with ONNX files")
    parser.add_argument("--int4", action="store_true", help="Also produce INT4 quantized models")
    parser.add_argument(
        "--protect-encoder-conv", action="store_true", default=True,
        help="Exclude encoder conv/batchnorm layers from INT8 quantization (recommended)"
    )
    args = parser.parse_args()

    encoder_path = os.path.join(args.input_dir, "cohere-encoder.onnx")
    decoder_path = os.path.join(args.input_dir, "cohere-decoder.onnx")

    # Nodes to exclude from encoder quantization:
    # The ConvSubsampling and BatchNorm layers are sensitive to quantization.
    # These node names may vary -- after export, inspect with:
    #   python -c "import onnx; m=onnx.load('...'); print([n.name for n in m.graph.node if 'Conv' in n.op_type or 'Batch' in n.op_type])"
    encoder_exclude = []
    if args.protect_encoder_conv:
        # Try to auto-detect conv/batchnorm node names
        try:
            import onnx
            model = onnx.load(encoder_path)
            for node in model.graph.node:
                if node.op_type in ("Conv", "BatchNormalization"):
                    encoder_exclude.append(node.name)
            if encoder_exclude:
                print(f"Protecting {len(encoder_exclude)} conv/batchnorm nodes from INT8 quantization")
            del model
        except Exception as e:
            print(f"Warning: Could not inspect encoder for conv nodes: {e}")

    # --- INT8 ---
    print("\n=== INT8 Dynamic Quantization ===")
    if os.path.exists(encoder_path):
        # Encoder is >2GB, must use external data format
        quantize_int8(
            encoder_path,
            os.path.join(args.input_dir, "cohere-encoder.int8.onnx"),
            nodes_to_exclude=encoder_exclude,
            use_external_data=True,
        )
    else:
        print(f"  Skipping encoder (not found: {encoder_path})")

    if os.path.exists(decoder_path):
        quantize_int8(
            decoder_path,
            os.path.join(args.input_dir, "cohere-decoder.int8.onnx"),
        )
    else:
        print(f"  Skipping decoder (not found: {decoder_path})")

    # --- INT4 ---
    if args.int4:
        print("\n=== INT4 Weight-Only Quantization ===")
        if os.path.exists(encoder_path):
            quantize_int4(
                encoder_path,
                os.path.join(args.input_dir, "cohere-encoder.int4.onnx"),
            )

        if os.path.exists(decoder_path):
            quantize_int4(
                decoder_path,
                os.path.join(args.input_dir, "cohere-decoder.int4.onnx"),
            )

    print("\n=== Done ===")
    print("Files in:", args.input_dir)
    for f in sorted(os.listdir(args.input_dir)):
        size = os.path.getsize(os.path.join(args.input_dir, f)) / (1024 * 1024)
        print(f"  {f:40s} {size:>8.1f} MB")


if __name__ == "__main__":
    main()