"""FP16 (half-precision) quantization for ONNX models.

Converts all eligible FP32 weights and activations to FP16, roughly halving
the on-disk model size and improving inference throughput on hardware with
native FP16 support (GPUs, Apple ANE, Qualcomm DSPs).
"""

from __future__ import annotations

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

logger = logging.getLogger(__name__)


def quantize_fp16(
    model_path: Union[str, Path],
    output_path: Union[str, Path, None] = None,
    keep_io_types: bool = True,
    min_positive_val: float = 1e-7,
    max_finite_val: float = 1e4,
) -> str:
    """Convert an ONNX model from FP32 to FP16 (half-precision).

    Uses ``onnxconverter_common.float16.convert_float_to_float16`` which
    rewrites the graph in-place, casting weights and intermediate tensors
    to ``float16``.

    Args:
        model_path: Path to the source FP32 ``.onnx`` model.
        output_path: Destination path for the FP16 model.  Defaults to
            ``<model_path stem>_fp16.onnx`` in the same directory.
        keep_io_types: If ``True`` (default), the model's input and output
            tensors remain FP32 so callers do not need to change their
            pre-/post-processing pipelines.
        min_positive_val: Minimum positive value representable in FP16.
            Values smaller than this are flushed to zero.
        max_finite_val: Maximum finite value representable in FP16.  Values
            larger than this are clipped.

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

    Raises:
        ImportError: If ``onnx`` or ``onnxconverter-common`` is not installed.
        FileNotFoundError: If *model_path* does not exist.
    """
    try:
        import onnx
    except ImportError as exc:
        raise ImportError(
            "FP16 quantization requires 'onnx'. "
            "Install it with:  pip install onnx"
        ) from exc

    try:
        from onnxconverter_common import float16
    except ImportError as exc:
        raise ImportError(
            "FP16 quantization requires 'onnxconverter-common'. "
            "Install it with:  pip install onnxconverter-common"
        ) from exc

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

    # Derive output path
    if output_path is None:
        output_path = model_path.with_name(
            model_path.stem + "_fp16" + 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(
        "FP16 quantization: loading %s (%.2f MB) ...",
        model_path,
        original_size_mb,
    )

    model = onnx.load(str(model_path))

    model_fp16 = float16.convert_float_to_float16(
        model,
        min_positive_val=min_positive_val,
        max_finite_val=max_finite_val,
        keep_io_types=keep_io_types,
    )

    onnx.save(model_fp16, str(output_path))

    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(
        "FP16 quantization complete: %s (%.2f MB -> %.2f MB, %.1f%% reduction)",
        output_path,
        original_size_mb,
        quantized_size_mb,
        reduction_pct,
    )

    return str(output_path)


def get_fp16_model_info(model_path: Union[str, Path]) -> dict[str, Any]:
    """Return summary information about an FP16-quantized ONNX model.

    Args:
        model_path: Path to the ``.onnx`` model.

    Returns:
        A dict with keys:

        * ``path`` -- absolute model path
        * ``size_mb`` -- file size in megabytes
        * ``opset_version`` -- ONNX opset version
        * ``ir_version`` -- ONNX IR version
        * ``num_nodes`` -- total number of graph nodes
    """
    try:
        import onnx
    except ImportError as exc:
        raise ImportError(
            "Model info 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}")

    model = onnx.load(str(model_path))
    size_mb = model_path.stat().st_size / (1024 * 1024)
    opset = model.opset_import[0].version if model.opset_import else None

    return {
        "path": str(model_path.resolve()),
        "size_mb": round(size_mb, 2),
        "opset_version": opset,
        "ir_version": model.ir_version,
        "num_nodes": len(model.graph.node),
    }
