"""Quantization tests — FP16, INT8 dynamic, model info, and benchmarking."""

from __future__ import annotations

import os

import numpy as np
import onnx
import pytest
import torch

from matrice_export.quantization import (
    get_fp16_model_info,
    quantize_int8,
)

# ------------------------------------------------------------------ #
# Helper: export a dummy model to ONNX for quantization tests
# ------------------------------------------------------------------ #


@pytest.fixture
def onnx_path(tmp_path):
    """Export a slightly larger model to ONNX for quantization tests.

    Uses a model large enough that file size > 0.0 MB when rounded to 2 decimals.
    Exports with dynamo=False (legacy exporter) for maximum compatibility with
    onnxruntime quantization, which requires clean shape inference.
    """
    model = torch.nn.Sequential(
        torch.nn.Linear(64, 128),
        torch.nn.ReLU(),
        torch.nn.Linear(128, 64),
        torch.nn.ReLU(),
        torch.nn.Linear(64, 10),
    )
    model.eval()
    sample = torch.randn(1, 64)
    path = str(tmp_path / "model.onnx")
    torch.onnx.export(
        model, sample, path,
        opset_version=17,
        dynamo=False,
        input_names=["input"],
        output_names=["output"],
    )
    return path


# ------------------------------------------------------------------ #
# 1. FP16 quantization
# ------------------------------------------------------------------ #


def _fp16_available():
    try:
        from onnxconverter_common import float16  # noqa: F401
        return True
    except ImportError:
        return False


class TestQuantizeFP16:
    @pytest.mark.skipif(not _fp16_available(), reason="onnxconverter-common not installed")
    def test_fp16_produces_file(self, onnx_path, tmp_path):
        """FP16 quantization produces a valid output file."""
        from matrice_export.quantization import quantize_fp16

        fp16_path = quantize_fp16(onnx_path, output_path=str(tmp_path / "model_fp16.onnx"))
        assert os.path.isfile(fp16_path)

    @pytest.mark.skipif(not _fp16_available(), reason="onnxconverter-common not installed")
    def test_fp16_is_valid_onnx(self, onnx_path, tmp_path):
        """FP16 model passes onnx.checker.check_model."""
        from matrice_export.quantization import quantize_fp16

        fp16_path = quantize_fp16(onnx_path, output_path=str(tmp_path / "model_fp16.onnx"))
        model = onnx.load(fp16_path)
        onnx.checker.check_model(model)

    @pytest.mark.skipif(not _fp16_available(), reason="onnxconverter-common not installed")
    def test_fp16_default_output_path(self, onnx_path):
        """FP16 with default output_path creates *_fp16.onnx next to original."""
        from matrice_export.quantization import quantize_fp16

        fp16_path = quantize_fp16(onnx_path)
        assert fp16_path.endswith("_fp16.onnx")
        assert os.path.isfile(fp16_path)
        # Cleanup
        os.unlink(fp16_path)


# ------------------------------------------------------------------ #
# 2. INT8 dynamic quantization
# ------------------------------------------------------------------ #


class TestQuantizeINT8:
    def test_int8_produces_file(self, onnx_path, tmp_path):
        """INT8 dynamic quantization produces a valid output file."""
        int8_path = quantize_int8(onnx_path, output_path=str(tmp_path / "model_int8.onnx"))
        assert os.path.isfile(int8_path)

    def test_int8_is_valid_onnx(self, onnx_path, tmp_path):
        """INT8 model passes onnx.checker.check_model."""
        int8_path = quantize_int8(onnx_path, output_path=str(tmp_path / "model_int8.onnx"))
        model = onnx.load(int8_path)
        onnx.checker.check_model(model)

    def test_int8_default_output_path(self, onnx_path):
        """INT8 with default output_path creates *_int8.onnx next to original."""
        int8_path = quantize_int8(onnx_path)
        assert int8_path.endswith("_int8.onnx")
        assert os.path.isfile(int8_path)
        os.unlink(int8_path)

    def test_int8_inference_runs(self, onnx_path, tmp_path):
        """INT8 model can be loaded and produces output via ONNX Runtime."""
        import onnxruntime as ort

        int8_path = quantize_int8(onnx_path, output_path=str(tmp_path / "model_int8.onnx"))
        sess = ort.InferenceSession(int8_path)
        input_name = sess.get_inputs()[0].name
        np_input = np.random.randn(1, 64).astype(np.float32)
        result = sess.run(None, {input_name: np_input})
        assert result is not None
        assert len(result) > 0
        assert result[0].shape[1] == 10  # model outputs 10 classes


# ------------------------------------------------------------------ #
# 3. FP16 model info
# ------------------------------------------------------------------ #


class TestFP16ModelInfo:
    def test_info_returns_expected_keys(self, onnx_path):
        """get_fp16_model_info returns a dict with the expected keys."""
        info = get_fp16_model_info(onnx_path)
        assert "path" in info
        assert "size_mb" in info
        assert "opset_version" in info
        assert "ir_version" in info
        assert "num_nodes" in info

    def test_info_size_positive(self, onnx_path):
        """Reported size_mb is positive."""
        info = get_fp16_model_info(onnx_path)
        assert info["size_mb"] > 0

    def test_info_num_nodes_positive(self, onnx_path):
        """The model has at least one graph node."""
        info = get_fp16_model_info(onnx_path)
        assert info["num_nodes"] > 0


# ------------------------------------------------------------------ #
# 4. Edge cases
# ------------------------------------------------------------------ #


class TestQuantizeEdgeCases:
    def test_int8_file_not_found(self, tmp_path):
        """quantize_int8 raises FileNotFoundError for non-existent file."""
        with pytest.raises(FileNotFoundError):
            quantize_int8(str(tmp_path / "nonexistent.onnx"))

    def test_fp16_info_file_not_found(self, tmp_path):
        """get_fp16_model_info raises FileNotFoundError for non-existent file."""
        with pytest.raises(FileNotFoundError):
            get_fp16_model_info(str(tmp_path / "nonexistent.onnx"))
