"""INT8 (8-bit integer) quantization for ONNX models.

Provides both *dynamic* and *static* INT8 quantization.  Dynamic quantization
quantizes weights offline and activations on-the-fly at inference time,
requiring no calibration data.  Static quantization pre-computes activation
ranges using a calibration dataset for tighter quantization and typically
better accuracy.
"""

from __future__ import annotations

import logging
from pathlib import Path
from typing import Any, Sequence, Union

logger = logging.getLogger(__name__)


# --------------------------------------------------------------------------- #
# Dynamic INT8 quantization
# --------------------------------------------------------------------------- #

def quantize_int8(
    model_path: Union[str, Path],
    output_path: Union[str, Path, None] = None,
    per_channel: bool = True,
    reduce_range: bool = False,
    nodes_to_exclude: list[str] | None = None,
) -> str:
    """Apply dynamic INT8 quantization to an ONNX model.

    Quantizes weights to 8-bit integers.  Activations are quantized
    dynamically at inference time.  Typical size reduction is 3-4x compared
    to FP32.

    Args:
        model_path: Path to the source ``.onnx`` model.
        output_path: Destination path for the quantized model.  Defaults to
            ``<model_path stem>_int8.onnx`` in the same directory.
        per_channel: If ``True``, apply per-channel quantization (more
            accurate but slightly slower on some hardware).
        reduce_range: If ``True``, use 7-bit ranges for activations.  May
            improve accuracy on older CPUs that lack VNNI.
        nodes_to_exclude: Optional list of node names to skip during
            quantization (e.g. sensitive layers).

    Returns:
        Absolute path to the saved INT8 model as a string.

    Raises:
        ImportError: If ``onnxruntime`` is not installed.
        FileNotFoundError: If *model_path* does not exist.
    """
    try:
        from onnxruntime.quantization import QuantType, quantize_dynamic
    except ImportError as exc:
        raise ImportError(
            "INT8 dynamic quantization requires 'onnxruntime'. "
            "Install it with:  pip install onnxruntime"
        ) from exc

    model_path = Path(model_path)
    if not model_path.exists():
        raise FileNotFoundError(f"ONNX model not found: {model_path}")

    if output_path is None:
        output_path = model_path.with_name(
            model_path.stem + "_int8" + model_path.suffix
        )
    output_path = Path(output_path)
    output_path.parent.mkdir(parents=True, exist_ok=True)

    original_size_mb = model_path.stat().st_size / (1024 * 1024)
    logger.info(
        "INT8 dynamic quantization: loading %s (%.2f MB) ...",
        model_path,
        original_size_mb,
    )

    extra_options: dict[str, Any] = {}
    if nodes_to_exclude:
        extra_options["nodes_to_exclude"] = nodes_to_exclude

    quantize_dynamic(
        model_input=str(model_path),
        model_output=str(output_path),
        weight_type=QuantType.QInt8,
        per_channel=per_channel,
        reduce_range=reduce_range,
        extra_options=extra_options if extra_options else None,
    )

    quantized_size_mb = output_path.stat().st_size / (1024 * 1024)
    reduction_pct = (
        (1.0 - quantized_size_mb / original_size_mb) * 100.0
        if original_size_mb > 0
        else 0.0
    )

    logger.info(
        "INT8 dynamic quantization complete: %s (%.2f MB -> %.2f MB, %.1f%% reduction)",
        output_path,
        original_size_mb,
        quantized_size_mb,
        reduction_pct,
    )

    return str(output_path)


# --------------------------------------------------------------------------- #
# Static INT8 quantization
# --------------------------------------------------------------------------- #

class _CalibrationDataReader:
    """Minimal calibration data reader for static INT8 quantization.

    Wraps an iterable of ``numpy.ndarray`` samples into the interface
    expected by ``onnxruntime.quantization.CalibrationDataReader``.
    """

    def __init__(self, input_name: str, data: Sequence[Any]) -> None:
        self._input_name = input_name
        self._data = iter(data)

    def get_next(self) -> dict[str, Any] | None:
        try:
            sample = next(self._data)
            return {self._input_name: sample}
        except StopIteration:
            return None


def quantize_int8_static(
    model_path: Union[str, Path],
    output_path: Union[str, Path, None] = None,
    calibration_data: Sequence[Any] | None = None,
    per_channel: bool = True,
    reduce_range: bool = False,
    nodes_to_exclude: list[str] | None = None,
) -> str:
    """Apply static INT8 quantization to an ONNX model.

    Static quantization pre-computes quantization parameters for both
    weights *and* activations using calibration data, typically yielding
    better accuracy than dynamic quantization at the cost of requiring
    representative inputs.

    Args:
        model_path: Path to the source ``.onnx`` model.
        output_path: Destination path for the quantized model.  Defaults to
            ``<model_path stem>_int8_static.onnx``.
        calibration_data: A sequence (list / tuple) of ``numpy.ndarray``
            samples used for calibration.  Each sample must match the
            model's first input shape and dtype.  If ``None``, falls back
            to dynamic quantization with a warning.
        per_channel: If ``True``, apply per-channel quantization.
        reduce_range: If ``True``, use 7-bit activation ranges.
        nodes_to_exclude: Optional list of node names to skip.

    Returns:
        Absolute path to the saved INT8 model as a string.

    Raises:
        ImportError: If ``onnxruntime`` is not installed.
        FileNotFoundError: If *model_path* does not exist.
    """
    if calibration_data is None or len(calibration_data) == 0:
        logger.warning(
            "No calibration data provided; falling back to dynamic INT8 "
            "quantization."
        )
        return quantize_int8(
            model_path,
            output_path=output_path,
            per_channel=per_channel,
            reduce_range=reduce_range,
            nodes_to_exclude=nodes_to_exclude,
        )

    try:
        from onnxruntime.quantization import QuantType, quantize_static
    except ImportError as exc:
        raise ImportError(
            "INT8 static quantization requires 'onnxruntime'. "
            "Install it with:  pip install onnxruntime"
        ) from exc

    try:
        import onnx
    except ImportError as exc:
        raise ImportError(
            "INT8 static quantization requires 'onnx'. "
            "Install it with:  pip install onnx"
        ) from exc

    model_path = Path(model_path)
    if not model_path.exists():
        raise FileNotFoundError(f"ONNX model not found: {model_path}")

    if output_path is None:
        output_path = model_path.with_name(
            model_path.stem + "_int8_static" + model_path.suffix
        )
    output_path = Path(output_path)
    output_path.parent.mkdir(parents=True, exist_ok=True)

    original_size_mb = model_path.stat().st_size / (1024 * 1024)
    logger.info(
        "INT8 static quantization: loading %s (%.2f MB) with %d calibration samples ...",
        model_path,
        original_size_mb,
        len(calibration_data),
    )

    # Determine input name from the model
    model_proto = onnx.load(str(model_path))
    input_name = model_proto.graph.input[0].name

    data_reader = _CalibrationDataReader(input_name, calibration_data)

    extra_options: dict[str, Any] = {}
    if nodes_to_exclude:
        extra_options["nodes_to_exclude"] = nodes_to_exclude

    quantize_static(
        model_input=str(model_path),
        model_output=str(output_path),
        calibration_data_reader=data_reader,
        weight_type=QuantType.QInt8,
        per_channel=per_channel,
        reduce_range=reduce_range,
        extra_options=extra_options if extra_options else None,
    )

    quantized_size_mb = output_path.stat().st_size / (1024 * 1024)
    reduction_pct = (
        (1.0 - quantized_size_mb / original_size_mb) * 100.0
        if original_size_mb > 0
        else 0.0
    )

    logger.info(
        "INT8 static quantization complete: %s (%.2f MB -> %.2f MB, %.1f%% reduction)",
        output_path,
        original_size_mb,
        quantized_size_mb,
        reduction_pct,
    )

    return str(output_path)
