"""Compute-device resolution for training and inference.

Extracted from the ``update_compute()`` functions duplicated across all
four BYOM codebases (classification, YOLOv10, D-FINE, r2plus1d).
"""

from __future__ import annotations

from typing import TYPE_CHECKING, Any

from matrice_models.common.device import move_to_device as common_move_to_device
from matrice_models.common.device import resolve_device as common_resolve_device


if TYPE_CHECKING:
    import torch

def resolve_device(preferred: str = "cuda") -> torch.device:
    """Select the best available compute device.

    Tries the *preferred* backend first, then falls back through the
    priority chain ``cuda → mps → cpu``.

    Args:
        preferred: Requested device — one of ``"cuda"``, ``"mps"``,
            or ``"cpu"``.

    Returns:
        A :class:`torch.device` for the selected backend.

    Example::

        device = resolve_device("cuda")
        model = model.to(device)
    """
    return common_resolve_device(preferred)


def move_to_device(model: Any, device: torch.device) -> torch.device:
    """Move a model to *device* and return the device used.

    Handles both plain ``nn.Module`` and wrapper objects that store the
    actual model in a ``.model`` attribute (e.g. the ``r2plus1d_Model``
    wrapper from the R2Plus1D BYOM repo).

    Args:
        model: A PyTorch ``nn.Module``, or a wrapper with a ``.model``
            attribute pointing to the actual module.
        device: Target device.

    Returns:
        The device the model was moved to (same as *device*).
    """
    return common_move_to_device(model, device)
