"""ONNX-specific export tests — file validity, output correctness, and export options."""

from __future__ import annotations

import os

import numpy as np
import onnx
import onnxruntime as ort
import pytest
import torch

from matrice_export import ExportPipeline

# ------------------------------------------------------------------ #
# 1. ONNX file passes onnx.checker
# ------------------------------------------------------------------ #

class TestOnnxValidity:
    def test_onnx_file_valid(self, dummy_model, sample_input, tmp_path):
        """The exported .onnx file passes onnx.checker.check_model without errors."""
        pipeline = ExportPipeline(dummy_model, sample_input, task="classification")
        results = pipeline.export(["onnx"], str(tmp_path), validate=False)
        model_onnx = onnx.load(results["onnx"]["path"])
        # This will raise if invalid
        onnx.checker.check_model(model_onnx)


# ------------------------------------------------------------------ #
# 2. ONNX output shape matches PyTorch output
# ------------------------------------------------------------------ #

class TestOnnxOutputShape:
    def test_output_shape_matches(self, dummy_model, sample_input, tmp_path):
        """The ONNX model output shape matches the original PyTorch model output shape."""
        dummy_model.eval()
        with torch.no_grad():
            pt_output = dummy_model(sample_input).numpy()

        pipeline = ExportPipeline(dummy_model, sample_input, task="classification")
        results = pipeline.export(["onnx"], str(tmp_path), validate=False)

        sess = ort.InferenceSession(results["onnx"]["path"])
        input_name = sess.get_inputs()[0].name
        onnx_output = sess.run(None, {input_name: sample_input.numpy()})[0]

        assert pt_output.shape == onnx_output.shape


# ------------------------------------------------------------------ #
# 3. ONNX output values close to PyTorch
# ------------------------------------------------------------------ #

class TestOnnxOutputValues:
    def test_output_values_close(self, dummy_model, sample_input, tmp_path):
        """The ONNX model output values are numerically close to PyTorch output."""
        dummy_model.eval()
        with torch.no_grad():
            pt_output = dummy_model(sample_input).numpy()

        pipeline = ExportPipeline(dummy_model, sample_input, task="classification")
        results = pipeline.export(["onnx"], str(tmp_path), validate=False)

        sess = ort.InferenceSession(results["onnx"]["path"])
        input_name = sess.get_inputs()[0].name
        onnx_output = sess.run(None, {input_name: sample_input.numpy()})[0]

        np.testing.assert_allclose(pt_output, onnx_output, rtol=1e-5, atol=1e-5)


# ------------------------------------------------------------------ #
# 4. ONNX with simplify=True (requires onnxsim; skip if missing)
# ------------------------------------------------------------------ #

def _onnxsim_available() -> bool:
    try:
        import onnxsim  # noqa: F401
        return True
    except ImportError:
        return False


class TestOnnxSimplify:
    @pytest.mark.skipif(
        not _onnxsim_available(),
        reason="onnx-simplifier not installed",
    )
    def test_onnx_simplify(self, dummy_model, sample_input, tmp_path):
        """Export with simplify=True produces a valid ONNX file."""
        pipeline = ExportPipeline(dummy_model, sample_input, task="classification")
        results = pipeline.export(["onnx"], str(tmp_path), validate=False, simplify=True)

        assert results["onnx"]["status"] == "success"
        model_onnx = onnx.load(results["onnx"]["path"])
        onnx.checker.check_model(model_onnx)


class TestOnnxSimplifyUnavailable:
    @pytest.mark.skipif(
        _onnxsim_available(),
        reason="onnx-simplifier IS installed; skip unavailable test",
    )
    def test_onnx_simplify_import_error(self, dummy_model, sample_input, tmp_path):
        """When onnxsim is not installed, simplify=True captures an ImportError."""
        pipeline = ExportPipeline(dummy_model, sample_input, task="classification")
        results = pipeline.export(["onnx"], str(tmp_path), validate=False, simplify=True)
        assert results["onnx"]["status"] == "error"
        assert "onnx-simplifier" in results["onnx"]["error"]


# ------------------------------------------------------------------ #
# 5. ONNX with dynamic=True
# ------------------------------------------------------------------ #

class TestOnnxDynamic:
    def test_onnx_dynamic_axes(self, conv_model, image_input, tmp_path):
        """Export with dynamic=True produces a valid model with dynamic batch dimension."""
        pipeline = ExportPipeline(conv_model, image_input, task="classification")
        results = pipeline.export(["onnx"], str(tmp_path), validate=False, dynamic=True)

        assert results["onnx"]["status"] == "success"

        model_onnx = onnx.load(results["onnx"]["path"])
        onnx.checker.check_model(model_onnx)

        # Check that the input has a dynamic dimension
        input_shape = model_onnx.graph.input[0].type.tensor_type.shape
        dim0 = input_shape.dim[0]
        # A dynamic dim has a dim_param (string name) rather than a fixed dim_value
        assert dim0.dim_param != "" or dim0.dim_value == 0, (
            "Expected dynamic batch dimension on input"
        )


# ------------------------------------------------------------------ #
# 6. ONNX metadata contains stride
# ------------------------------------------------------------------ #

class TestOnnxMetadata:
    def test_metadata_contains_stride(self, dummy_model, sample_input, tmp_path):
        """The exported ONNX file embeds a 'stride' metadata property."""
        pipeline = ExportPipeline(dummy_model, sample_input, task="classification")
        results = pipeline.export(["onnx"], str(tmp_path), validate=False)

        model_onnx = onnx.load(results["onnx"]["path"])
        metadata_keys = [prop.key for prop in model_onnx.metadata_props]
        assert "stride" in metadata_keys

    def test_metadata_contains_names(self, dummy_model, sample_input, tmp_path):
        """The exported ONNX file embeds a 'names' metadata property."""
        pipeline = ExportPipeline(dummy_model, sample_input, task="classification")
        results = pipeline.export(["onnx"], str(tmp_path), validate=False)

        model_onnx = onnx.load(results["onnx"]["path"])
        metadata_keys = [prop.key for prop in model_onnx.metadata_props]
        assert "names" in metadata_keys


# ------------------------------------------------------------------ #
# 7. ONNX file size is reasonable
# ------------------------------------------------------------------ #

class TestOnnxFileSize:
    def test_file_size_under_10mb(self, dummy_model, sample_input, tmp_path):
        """A small model exports to an ONNX file well under 10 MB."""
        pipeline = ExportPipeline(dummy_model, sample_input, task="classification")
        results = pipeline.export(["onnx"], str(tmp_path), validate=False)

        size_bytes = os.path.getsize(results["onnx"]["path"])
        assert size_bytes < 10 * 1024 * 1024, "ONNX file exceeds 10 MB for a trivial model"

    def test_file_size_positive(self, dummy_model, sample_input, tmp_path):
        """The exported ONNX file has a positive size (not empty)."""
        pipeline = ExportPipeline(dummy_model, sample_input, task="classification")
        results = pipeline.export(["onnx"], str(tmp_path), validate=False)

        size_bytes = os.path.getsize(results["onnx"]["path"])
        assert size_bytes > 0


# ------------------------------------------------------------------ #
# 8. ONNX export with conv model (image input)
# ------------------------------------------------------------------ #

class TestOnnxConvModel:
    def test_conv_model_export(self, conv_model, image_input, tmp_path):
        """A CNN model exports successfully to ONNX."""
        pipeline = ExportPipeline(conv_model, image_input, task="classification")
        results = pipeline.export(["onnx"], str(tmp_path), validate=False)

        assert results["onnx"]["status"] == "success"
        model_onnx = onnx.load(results["onnx"]["path"])
        onnx.checker.check_model(model_onnx)

    def test_conv_model_output_values_close(self, conv_model, image_input, tmp_path):
        """ONNX inference on a CNN model matches PyTorch output numerically."""
        conv_model.eval()
        with torch.no_grad():
            pt_output = conv_model(image_input).numpy()

        pipeline = ExportPipeline(conv_model, image_input, task="classification")
        results = pipeline.export(["onnx"], str(tmp_path), validate=False)

        sess = ort.InferenceSession(results["onnx"]["path"])
        input_name = sess.get_inputs()[0].name
        onnx_output = sess.run(None, {input_name: image_input.numpy()})[0]

        np.testing.assert_allclose(pt_output, onnx_output, rtol=1e-4, atol=1e-4)
