"""Device-agnostic (CPU/CUDA) utilities (D4).

Designed so the prototype pipeline runs on the current CPU-only cluster and moves
to GPU unchanged. Imports torch LAZILY -- the module loads even though torch is
not yet installed, so metadata-independent scaffolding and tests are unaffected.
"""
from __future__ import annotations


def torch_available():
    try:
        import torch  # noqa: F401
        return True
    except Exception:
        return False


def get_device(prefer="auto"):
    """Return a device string ('cpu'/'cuda'). If torch is present, returns a
    torch.device instead. ``prefer`` in {auto, cpu, cuda}."""
    if not torch_available():
        return "cpu"
    import torch
    if prefer == "cpu":
        return torch.device("cpu")
    if prefer in ("auto", "cuda") and torch.cuda.is_available():
        return torch.device("cuda")
    return torch.device("cpu")


def device_report():
    """Human-readable device/availability summary for logs and the run record."""
    if not torch_available():
        return {"torch": False, "cuda": False, "device": "cpu",
                "note": "torch not installed; CPU-only prototype"}
    import torch
    cuda = torch.cuda.is_available()
    r = {"torch": torch.__version__, "cuda": cuda,
         "device": "cuda" if cuda else "cpu"}
    if cuda:
        r["gpus"] = [torch.cuda.get_device_name(i) for i in range(torch.cuda.device_count())]
    return r


def estimate_tensor_memory_mb(shape, dtype_bytes=4):
    """Approximate memory (MB) for a dense tensor of ``shape``. Metadata-free
    helper for sizing the CPU prototype before any model exists."""
    n = 1
    for s in shape:
        n *= int(s)
    return n * dtype_bytes / (1024 ** 2)


def autocast(enabled=True):
    """Mixed-precision context manager if torch is present, else a no-op.
    Only meaningful on GPU; harmless on CPU."""
    if not torch_available():
        from contextlib import nullcontext
        return nullcontext()
    import torch
    dev = "cuda" if torch.cuda.is_available() else "cpu"
    return torch.autocast(device_type=dev, enabled=enabled and dev == "cuda")