"""MPS (Apple Silicon GPU) inference utilities.

Provides helpers for loading PyTorch models onto the Metal Performance
Shaders (MPS) device, running inference, and benchmarking latency.
Falls back to CPU transparently when MPS is not available.
"""

from __future__ import annotations

import logging
import time
from pathlib import Path
from typing import Any, Union

logger = logging.getLogger(__name__)


# --------------------------------------------------------------------------- #
# Device helpers
# --------------------------------------------------------------------------- #

def _get_mps_device() -> Any:
    """Return the MPS device if available, otherwise CPU.

    Returns:
        ``torch.device("mps")`` when MPS is available and functional,
        ``torch.device("cpu")`` otherwise.
    """
    try:
        import torch
    except ImportError as exc:
        raise ImportError(
            "MPS inference requires 'torch'. "
            "Install it with:  pip install torch"
        ) from exc

    if hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
        return torch.device("mps")

    logger.info("MPS not available; falling back to CPU.")
    return torch.device("cpu")


# --------------------------------------------------------------------------- #
# Load
# --------------------------------------------------------------------------- #

def load_model_mps(
    model_path: Union[str, Path],
) -> tuple[Any, Any]:
    """Load a PyTorch model onto the MPS device (Apple Silicon GPU).

    Falls back to CPU if MPS is not available.  The model is placed in
    ``eval`` mode automatically.

    Args:
        model_path: File-system path to a saved PyTorch model (``.pt`` /
            ``.pth`` file saved via ``torch.save``).

    Returns:
        A ``(model, device)`` tuple where *model* is the loaded
        ``torch.nn.Module`` and *device* is the ``torch.device`` it
        resides on (``"mps"`` or ``"cpu"``).

    Raises:
        ImportError: If ``torch`` is not installed.
        FileNotFoundError: If *model_path* does not exist.
    """
    try:
        import torch
    except ImportError as exc:
        raise ImportError(
            "MPS inference requires 'torch'. "
            "Install it with:  pip install torch"
        ) from exc

    model_path = Path(model_path)
    if not model_path.exists():
        raise FileNotFoundError(f"Model file not found: {model_path}")

    device = _get_mps_device()
    logger.info("Loading model from %s onto %s ...", model_path, device)

    model = torch.load(str(model_path), map_location=device, weights_only=False)
    model.eval()
    return model, device


# --------------------------------------------------------------------------- #
# Inference
# --------------------------------------------------------------------------- #

def run_inference_mps(
    model: Any,
    input_tensor: Any,
    device: Any = None,
) -> Any:
    """Run inference on the MPS device (or whichever device is provided).

    Args:
        model: A PyTorch ``nn.Module`` already on the target device.
        input_tensor: Input ``torch.Tensor``.  Will be moved to *device*
            if it is not already there.
        device: ``torch.device`` to use.  If ``None``, the device is
            inferred from the model's first parameter.

    Returns:
        Model output tensor(s), moved back to CPU for downstream
        consumption.

    Raises:
        ImportError: If ``torch`` is not installed.
    """
    try:
        import torch
    except ImportError as exc:
        raise ImportError(
            "MPS inference requires 'torch'. "
            "Install it with:  pip install torch"
        ) from exc

    if device is None:
        try:
            device = next(model.parameters()).device
        except StopIteration:
            device = torch.device("cpu")

    input_tensor = input_tensor.to(device)

    with torch.no_grad():
        output = model(input_tensor)

    # Move outputs back to CPU for interoperability
    if isinstance(output, torch.Tensor):
        return output.cpu()
    if isinstance(output, (tuple, list)):
        return type(output)(
            o.cpu() if isinstance(o, torch.Tensor) else o for o in output
        )
    return output


# --------------------------------------------------------------------------- #
# Benchmark
# --------------------------------------------------------------------------- #

def benchmark_mps(
    model: Any,
    input_tensor: Any,
    n_runs: int = 100,
    warmup_runs: int = 10,
) -> dict[str, Any]:
    """Benchmark inference latency on the MPS device.

    Performs a configurable number of warm-up iterations (discarded) followed
    by *n_runs* timed iterations.  Reports mean, median, min, and max
    latency in milliseconds.

    Args:
        model: A PyTorch ``nn.Module`` already on the target device.
        input_tensor: Input ``torch.Tensor``.
        n_runs: Number of timed inference runs.
        warmup_runs: Number of warm-up runs (not timed).

    Returns:
        A dict with keys:

        * ``device_name`` -- string representation of the device
        * ``n_runs`` -- number of timed runs
        * ``latency_mean_ms`` -- mean latency in milliseconds
        * ``latency_median_ms`` -- median latency in milliseconds
        * ``latency_min_ms`` -- minimum latency
        * ``latency_max_ms`` -- maximum latency
        * ``throughput_fps`` -- estimated frames per second

    Raises:
        ImportError: If ``torch`` or ``numpy`` is not installed.
    """
    try:
        import torch
    except ImportError as exc:
        raise ImportError(
            "MPS benchmarking requires 'torch'. "
            "Install it with:  pip install torch"
        ) from exc

    try:
        import numpy as np
    except ImportError as exc:
        raise ImportError(
            "MPS benchmarking requires 'numpy'. "
            "Install it with:  pip install numpy"
        ) from exc

    # Determine device
    try:
        device = next(model.parameters()).device
    except StopIteration:
        device = torch.device("cpu")

    input_tensor = input_tensor.to(device)
    device_name = str(device)
    logger.info(
        "Benchmarking on %s: %d warmup + %d timed runs ...",
        device_name,
        warmup_runs,
        n_runs,
    )

    model.eval()

    # Warm-up
    with torch.no_grad():
        for _ in range(warmup_runs):
            _ = model(input_tensor)
            # Synchronise MPS to ensure work is complete before timing
            if device.type == "mps":
                torch.mps.synchronize()

    # Timed runs
    latencies: list[float] = []
    with torch.no_grad():
        for _ in range(n_runs):
            t_start = time.perf_counter()
            _ = model(input_tensor)
            if device.type == "mps":
                torch.mps.synchronize()
            t_end = time.perf_counter()
            latencies.append((t_end - t_start) * 1000.0)  # ms

    latencies_np = np.array(latencies)
    mean_ms = float(np.mean(latencies_np))

    result: dict[str, Any] = {
        "device_name": device_name,
        "n_runs": n_runs,
        "latency_mean_ms": mean_ms,
        "latency_median_ms": float(np.median(latencies_np)),
        "latency_min_ms": float(np.min(latencies_np)),
        "latency_max_ms": float(np.max(latencies_np)),
        "throughput_fps": 1000.0 / mean_ms if mean_ms > 0 else 0.0,
    }

    logger.info(
        "Benchmark results on %s: mean=%.2f ms, median=%.2f ms, "
        "min=%.2f ms, max=%.2f ms, throughput=%.1f FPS",
        device_name,
        result["latency_mean_ms"],
        result["latency_median_ms"],
        result["latency_min_ms"],
        result["latency_max_ms"],
        result["throughput_fps"],
    )

    return result
