"""Device selection and reproducibility helpers.

Centralised so the same logic can't drift between trainers and inference.
"""

from __future__ import annotations

import random

import numpy as np
import torch


def pick_device() -> torch.device:
    """Return the most capable torch device available.

    Order of preference: Apple Silicon MPS, NVIDIA CUDA, CPU. Chosen here
    because the project's primary target is M-series Macs; CUDA support is
    a free bonus when running on a workstation.
    """
    if torch.backends.mps.is_available():
        return torch.device("mps")
    if torch.cuda.is_available():
        return torch.device("cuda")
    return torch.device("cpu")


def seed_everything(seed: int) -> None:
    """Set Python, NumPy and PyTorch RNG seeds for reproducible runs.

    Note: full bitwise reproducibility on MPS isn't guaranteed by PyTorch;
    this gets us *seed-driven* identical splits and broadly identical training
    trajectories, which is what we want for science.
    """
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    if torch.cuda.is_available():
        torch.cuda.manual_seed_all(seed)