"""TorchScript validator for exported ``.torchscript`` models."""

from __future__ import annotations

import logging
from typing import TYPE_CHECKING, Any

if TYPE_CHECKING:
    import numpy as np

from matrice_export.validators.base import BaseValidator

logger = logging.getLogger(__name__)


class TorchScriptValidator(BaseValidator):
    """Validate an exported TorchScript model.

    Loads the model with ``torch.jit.load``, runs inference on the
    provided *sample_input* (converted to a ``torch.Tensor``), compares
    the output against an optional *baseline*, and benchmarks median
    inference latency.
    """

    def validate(
        self,
        model_path: str,
        sample_input: np.ndarray,
        baseline: np.ndarray | None = None,
        atol: float = 1e-5,
        rtol: float = 1e-3,
        **kwargs: Any,
    ) -> dict[str, Any]:
        """Validate a TorchScript model.

        Parameters
        ----------
        model_path:
            Path to the ``.torchscript`` (or ``.pt``) file.
        sample_input:
            Numpy array (typically ``float32``, NCHW layout).
        baseline:
            Optional reference output from the original PyTorch model.
        atol:
            Absolute tolerance for value comparison.
        rtol:
            Relative tolerance for value comparison.
        **kwargs:
            Extra arguments (ignored).

        Returns
        -------
        dict
            Validation result with keys ``shape_match``, ``values_match``,
            ``max_diff``, ``latency_ms``, ``output_shape``.
        """
        try:
            import torch
        except ImportError:
            logger.warning(
                "PyTorch is not installed -- skipping TorchScript validation. "
                "Install it with:  pip install torch"
            )
            return {
                "status": "skipped",
                "reason": "torch is not installed",
            }

        logger.info("TorchScriptValidator: loading %s", model_path)

        # ------------------------------------------------------------------
        # 1. Load the TorchScript model
        # ------------------------------------------------------------------
        try:
            device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
            model = torch.jit.load(model_path, map_location=device)
            model.eval()
        except Exception as exc:
            logger.error("TorchScriptValidator: failed to load model: %s", exc)
            return {"status": "error", "error": str(exc)}

        # ------------------------------------------------------------------
        # 2. Prepare input tensor and run inference
        # ------------------------------------------------------------------
        try:
            tensor_input = torch.from_numpy(sample_input).to(device)
            with torch.no_grad():
                raw_output = model(tensor_input)

            # Handle tuple/list outputs (take first element)
            if isinstance(raw_output, (tuple, list)):
                raw_output = raw_output[0]

            output = raw_output.detach().cpu().numpy()
        except Exception as exc:
            logger.error("TorchScriptValidator: inference failed: %s", exc)
            return {"status": "error", "error": str(exc)}

        # ------------------------------------------------------------------
        # 3 & 4. Compare shapes and values against baseline
        # ------------------------------------------------------------------
        comparison = self._compare_outputs(output, baseline, atol=atol, rtol=rtol)

        # ------------------------------------------------------------------
        # 5. Benchmark latency (median of 100 runs)
        # ------------------------------------------------------------------
        def _run() -> None:
            with torch.no_grad():
                model(tensor_input)

        latency_ms = self._benchmark_latency(_run)

        logger.info(
            "TorchScriptValidator: shape_match=%s  values_match=%s  max_diff=%s  latency=%.2f ms",
            comparison["shape_match"],
            comparison["values_match"],
            comparison["max_diff"],
            latency_ms,
        )

        return {
            **comparison,
            "latency_ms": latency_ms,
        }
