from contextlib import nullcontext

import torch


def mps_is_available() -> bool:
    return hasattr(torch.backends, "mps") and torch.backends.mps.is_available()


def get_accelerator_device() -> torch.device:
    if torch.cuda.is_available():
        return torch.device("cuda")
    if mps_is_available():
        return torch.device("mps")
    return torch.device("cpu")


def is_accelerator_device(device) -> bool:
    if device is None:
        return False
    return torch.device(device).type in {"cuda", "mps"}


def accelerator_autocast(dtype=torch.bfloat16):
    device_type = get_accelerator_device().type
    if device_type in {"cuda", "mps"}:
        return torch.autocast(device_type=device_type, dtype=dtype)
    return nullcontext()


def empty_accelerator_cache():
    if torch.cuda.is_available():
        torch.cuda.synchronize()
        torch.cuda.empty_cache()
        try:
            torch.cuda.ipc_collect()
        except Exception:
            pass
    elif mps_is_available():
        torch.mps.synchronize()
        torch.mps.empty_cache()