"""Base validator interface for all format validators."""

from __future__ import annotations

import logging
import time
from abc import ABC, abstractmethod
from statistics import median
from typing import TYPE_CHECKING, Any

if TYPE_CHECKING:
    import numpy as np

logger = logging.getLogger(__name__)


class BaseValidator(ABC):
    """Abstract base class for post-export model validators.

    Every concrete validator (ONNX, TorchScript, OpenVINO, TensorRT, etc.)
    inherits from this class and implements :meth:`validate`.

    The validation contract:

    * Load the exported artefact at *model_path*.
    * Run inference on *sample_input*.
    * Optionally compare the output against a *baseline* (the original
      PyTorch model's output) using element-wise tolerances.
    * Benchmark median inference latency.
    * Return a result dictionary.
    """

    @abstractmethod
    def validate(
        self,
        model_path: str,
        sample_input: np.ndarray,
        baseline: np.ndarray | None = None,
        **kwargs: Any,
    ) -> dict[str, Any]:
        """Validate an exported model.

        Parameters
        ----------
        model_path:
            Filesystem path to the exported model artefact (file or
            directory, depending on the format).
        sample_input:
            A numpy array to feed through the model.  Typically NCHW
            ``float32``.
        baseline:
            Optional numpy array produced by the original PyTorch model.
            When provided, the validator checks shape and value agreement.
        **kwargs:
            Format-specific options (``atol``, ``rtol``, ``task``, etc.).

        Returns
        -------
        dict
            A dictionary with at least the following keys:

            * ``shape_match`` (bool) -- whether the output shape equals the
              baseline shape (``None`` if no baseline).
            * ``values_match`` (bool) -- whether values are close within
              tolerance (``None`` if no baseline).
            * ``max_diff`` (float) -- maximum absolute element-wise
              difference (``None`` if no baseline).
            * ``latency_ms`` (float) -- median inference latency in
              milliseconds over multiple runs.
            * ``output_shape`` (tuple) -- shape of the model output.
        """

    # ------------------------------------------------------------------
    # Shared helpers
    # ------------------------------------------------------------------
    @staticmethod
    def _benchmark_latency(
        run_fn: Any,
        n_runs: int = 100,
        warmup: int = 5,
    ) -> float:
        """Run *run_fn* repeatedly and return the **median** latency in ms.

        Parameters
        ----------
        run_fn:
            A zero-argument callable that performs a single inference.
        n_runs:
            Number of timed iterations (default ``100``).
        warmup:
            Number of untimed warm-up iterations (default ``5``).

        Returns
        -------
        float
            Median wall-clock latency in milliseconds.
        """
        # Warm-up (allows JIT compilers / caches to settle)
        for _ in range(warmup):
            run_fn()

        timings: list[float] = []
        for _ in range(n_runs):
            t0 = time.perf_counter()
            run_fn()
            t1 = time.perf_counter()
            timings.append((t1 - t0) * 1000.0)  # seconds -> ms

        return median(timings)

    @staticmethod
    def _compare_outputs(
        output: np.ndarray,
        baseline: np.ndarray | None,
        atol: float = 1e-5,
        rtol: float = 1e-3,
    ) -> dict[str, Any]:
        """Compare *output* against *baseline* and return match metadata.

        Parameters
        ----------
        output:
            The inference output from the exported model.
        baseline:
            The reference output from the original model (may be ``None``).
        atol:
            Absolute tolerance for ``np.allclose``.
        rtol:
            Relative tolerance for ``np.allclose``.

        Returns
        -------
        dict
            Keys: ``shape_match``, ``values_match``, ``max_diff``,
            ``output_shape``.
        """
        import numpy as np

        output_shape = tuple(output.shape)

        if baseline is None:
            return {
                "shape_match": None,
                "values_match": None,
                "max_diff": None,
                "output_shape": output_shape,
            }

        shape_match = output.shape == baseline.shape
        if shape_match:
            max_diff = float(np.max(np.abs(output - baseline)))
            values_match = bool(np.allclose(output, baseline, atol=atol, rtol=rtol))
        else:
            max_diff = float("inf")
            values_match = False

        return {
            "shape_match": shape_match,
            "values_match": values_match,
            "max_diff": max_diff,
            "output_shape": output_shape,
        }
