from typing import Optional

_TORCH = None
_TORCH_IMPORT_ERR: Optional[BaseException] = None
try:
    import torch as _TORCH
except Exception as e:
    _TORCH_IMPORT_ERR = e


def torch_available() -> bool:
    return _TORCH is not None


def ensure_torch():
    """Return torch module or raise with a platform-aware message."""
    if _TORCH is None:
        msg = (
            "PyTorch is not installed.\n"
            "- Apple Silicon (M-Series): `pip install torch`\n"
            "- macOS Intel: `pip install 'torch==2.2.*'` (or `conda install pytorch -c pytorch`)\n"
            "- Linux CPU: `pip install --index-url https://download.pytorch.org/whl/cpu torch`\n"
        )
        raise RuntimeError(msg) from _TORCH_IMPORT_ERR
    return _TORCH


class _NoGrad:
    def __call__(self, fn):  # decorator use: @no_grad()
        return fn

    def __enter__(self):  # context manager use: with no_grad():
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        return False


def no_grad():
    """Use as @no_grad() or `with no_grad():`."""
    if _TORCH is None:
        return _NoGrad()
    return _TORCH.no_grad()
