"""Benchmark quantization tests -- helpers and full benchmark_quantization()."""

from __future__ import annotations

import os

import numpy as np
import pytest
import torch

from matrice_export.quantization import quantize_int8
from matrice_export.quantization.benchmark import (
    _benchmark_latency,
    _compute_accuracy_delta,
    _get_file_size_mb,
    _run_ort_inference,
    benchmark_quantization,
)

# ------------------------------------------------------------------ #
# Local fixture: slightly larger model suitable for benchmarking
# ------------------------------------------------------------------ #


@pytest.fixture
def onnx_path_for_benchmark(tmp_path):
    """Export a model suitable for benchmarking (64->128->64->10)."""
    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 / "bench_model.onnx")
    torch.onnx.export(
        model,
        sample,
        path,
        opset_version=17,
        dynamo=False,
        input_names=["input"],
        output_names=["output"],
    )
    return path, sample.numpy()


# ------------------------------------------------------------------ #
# 1. Benchmark helper functions
# ------------------------------------------------------------------ #


class TestBenchmarkHelpers:
    """Tests for the internal benchmark helper functions."""

    def test_get_file_size_mb(self, onnx_path_for_benchmark):
        """_get_file_size_mb returns a positive float for a known ONNX file."""
        path, _ = onnx_path_for_benchmark
        size = _get_file_size_mb(path)
        assert isinstance(size, float)
        assert size > 0

    def test_run_ort_inference(self, onnx_path_for_benchmark):
        """_run_ort_inference returns a numpy array with the correct output shape."""
        path, sample_np = onnx_path_for_benchmark
        output = _run_ort_inference(path, sample_np)
        assert isinstance(output, np.ndarray)
        # Model output shape is (1, 10)
        assert output.shape == (1, 10)

    def test_benchmark_latency_returns_expected_keys(self, onnx_path_for_benchmark):
        """_benchmark_latency returns a dict with mean_ms, median_ms, min_ms, max_ms."""
        path, sample_np = onnx_path_for_benchmark
        result = _benchmark_latency(path, sample_np, n_runs=5, warmup_runs=1)
        assert isinstance(result, dict)
        for key in ("mean_ms", "median_ms", "min_ms", "max_ms"):
            assert key in result, f"Missing key: {key}"

    def test_benchmark_latency_values_positive(self, onnx_path_for_benchmark):
        """All latency values are strictly positive."""
        path, sample_np = onnx_path_for_benchmark
        result = _benchmark_latency(path, sample_np, n_runs=5, warmup_runs=1)
        for key in ("mean_ms", "median_ms", "min_ms", "max_ms"):
            assert result[key] > 0, f"{key} should be > 0, got {result[key]}"

    def test_compute_accuracy_delta_identical(self):
        """Identical arrays yield max_abs_diff~0 and cosine_similarity~1.0."""
        arr = np.random.randn(1, 10).astype(np.float32)
        delta = _compute_accuracy_delta(arr, arr)
        assert delta["max_abs_diff"] == pytest.approx(0.0, abs=1e-9)
        assert delta["mean_abs_diff"] == pytest.approx(0.0, abs=1e-9)
        assert delta["cosine_similarity"] == pytest.approx(1.0, abs=1e-6)

    def test_compute_accuracy_delta_different(self):
        """Different arrays have max_abs_diff > 0 and 0 < cosine_similarity < 1."""
        a = np.array([[1.0, 0.0, 0.0, 0.0, 0.0]], dtype=np.float32)
        b = np.array([[0.0, 1.0, 0.0, 0.0, 0.0]], dtype=np.float32)
        delta = _compute_accuracy_delta(a, b)
        assert delta["max_abs_diff"] > 0
        assert 0 <= delta["cosine_similarity"] < 1


# ------------------------------------------------------------------ #
# 2. Full benchmark_quantization() function
# ------------------------------------------------------------------ #


class TestBenchmarkQuantization:
    """Tests for the public benchmark_quantization() function."""

    def test_original_only_no_sample(self, onnx_path_for_benchmark):
        """With no sample_input, report has 'original' key with size_mb > 0 and latency is None."""
        path, _ = onnx_path_for_benchmark
        report = benchmark_quantization(path)
        assert "original" in report
        assert report["original"]["size_mb"] > 0
        assert report["original"]["latency"] is None
        assert report["original"]["accuracy"] is None

    def test_original_with_sample_input(self, onnx_path_for_benchmark):
        """With sample_input, latency is a dict and accuracy is a dict."""
        path, sample_np = onnx_path_for_benchmark
        report = benchmark_quantization(
            path, sample_input=sample_np, n_runs=5, warmup_runs=1
        )
        assert "original" in report
        assert isinstance(report["original"]["latency"], dict)
        assert "mean_ms" in report["original"]["latency"]
        assert isinstance(report["original"]["accuracy"], dict)

    def test_original_and_int8(self, onnx_path_for_benchmark, tmp_path):
        """Quantize to INT8, then benchmark both original and INT8 variants."""
        path, sample_np = onnx_path_for_benchmark
        int8_path = quantize_int8(path, output_path=str(tmp_path / "model_int8.onnx"))

        report = benchmark_quantization(
            path,
            int8_path=int8_path,
            sample_input=sample_np,
            n_runs=5,
            warmup_runs=1,
        )

        assert "original" in report
        assert "int8" in report
        assert report["original"]["size_mb"] > 0
        assert report["int8"]["size_mb"] > 0
        assert isinstance(report["int8"]["latency"], dict)

    def test_file_not_found_raises(self, tmp_path):
        """Non-existent original_path raises FileNotFoundError."""
        with pytest.raises(FileNotFoundError):
            benchmark_quantization(str(tmp_path / "nonexistent.onnx"))

    def test_missing_variant_path_skips(self, onnx_path_for_benchmark, tmp_path):
        """An fp16_path that does not exist is skipped (only 'original' in report)."""
        path, sample_np = onnx_path_for_benchmark
        fake_fp16_path = str(tmp_path / "does_not_exist_fp16.onnx")

        report = benchmark_quantization(
            path,
            fp16_path=fake_fp16_path,
            sample_input=sample_np,
            n_runs=5,
            warmup_runs=1,
        )

        assert "original" in report
        assert "fp16" not in report

    def test_pytorch_baseline_used(self, onnx_path_for_benchmark):
        """When pytorch_baseline is provided, accuracy is computed against it."""
        path, sample_np = onnx_path_for_benchmark
        # Run the model once to get a baseline
        baseline = _run_ort_inference(path, sample_np)

        report = benchmark_quantization(
            path,
            sample_input=sample_np,
            pytorch_baseline=baseline,
            n_runs=5,
            warmup_runs=1,
        )

        assert "original" in report
        accuracy = report["original"]["accuracy"]
        assert accuracy is not None
        assert "cosine_similarity" in accuracy
        # Since baseline == original output, accuracy should be near-perfect
        assert accuracy["cosine_similarity"] == pytest.approx(1.0, abs=1e-5)

    def test_report_structure(self, onnx_path_for_benchmark):
        """Each variant entry has path, size_mb, latency, and accuracy keys."""
        path, sample_np = onnx_path_for_benchmark
        report = benchmark_quantization(
            path, sample_input=sample_np, n_runs=5, warmup_runs=1
        )

        for variant_name, entry in report.items():
            assert "path" in entry, f"Missing 'path' in {variant_name}"
            assert "size_mb" in entry, f"Missing 'size_mb' in {variant_name}"
            assert "latency" in entry, f"Missing 'latency' in {variant_name}"
            assert "accuracy" in entry, f"Missing 'accuracy' in {variant_name}"

    def test_full_benchmark_with_int8(self, onnx_path_for_benchmark, tmp_path):
        """Full end-to-end: export model, quantize INT8, benchmark, verify report completeness."""
        path, sample_np = onnx_path_for_benchmark

        # Quantize to INT8
        int8_path = quantize_int8(path, output_path=str(tmp_path / "model_int8.onnx"))
        assert os.path.isfile(int8_path)

        # Run benchmark
        report = benchmark_quantization(
            path,
            int8_path=int8_path,
            sample_input=sample_np,
            n_runs=5,
            warmup_runs=1,
        )

        # Verify original
        assert report["original"]["size_mb"] > 0
        assert report["original"]["latency"] is not None
        assert report["original"]["latency"]["mean_ms"] > 0
        assert report["original"]["accuracy"] is not None

        # Verify INT8
        assert report["int8"]["size_mb"] > 0
        assert report["int8"]["latency"] is not None
        assert report["int8"]["latency"]["mean_ms"] > 0
        assert report["int8"]["accuracy"] is not None
        assert "cosine_similarity" in report["int8"]["accuracy"]
        assert "max_abs_diff" in report["int8"]["accuracy"]
        assert "mean_abs_diff" in report["int8"]["accuracy"]
