# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/60_bend.ipynb.

# %% auto 0
__all__ = ['isavail', 'isbuilt', 'getbackend', 'hasbackend', 'backend', 'hasmps', 'hascuda']

# %% ../nbs/60_bend.ipynb 6
from functools import wraps

# %% ../nbs/60_bend.ipynb 8
from types import ModuleType
from typing import Callable

# %% ../nbs/60_bend.ipynb 11
#| export

# %% ../nbs/60_bend.ipynb 13
#| export


# %% ../nbs/60_bend.ipynb 15
try: import torch
except ImportError: ...

# %% ../nbs/60_bend.ipynb 17
#| export


# %% ../nbs/60_bend.ipynb 19
from nlit import IS_AVAILABLE, IS_BUILT, MPS, CUDA
from nlit import __DOC__, __NAME__, __MODULE__, __QUALNAME__

# %% ../nbs/60_bend.ipynb 21
from .atyp import P

# %% ../nbs/60_bend.ipynb 24
def isavail(backend: ModuleType) -> bool:
    '''Check if the given backend is available.

    Parameters
    ----------
    backend : ModuleType
        The backend module to check for availability.

    Returns
    -------
    bool
        True if the backend is available, False otherwise.
    '''
    is_avail = getattr(backend, IS_AVAILABLE, lambda : False)
    return is_avail()

def isbuilt(backend: ModuleType) -> bool:
    '''Check if the given backend is built.

    Parameters
    ----------
    backend : ModuleType
        The backend module to check for build status.

    Returns
    -------
    bool
        True if the backend is built, False otherwise.
    '''
    is_built = getattr(backend, IS_BUILT, lambda : False)
    return is_built()

def getbackend(name: str) -> ModuleType | None:
    '''Retrieve a backend module by name.

    Parameters
    ----------
    name : str
        The name of the backend to retrieve.

    Returns
    -------
    ModuleType | None
        The backend module if found, None otherwise.
    '''
    if isinstance(name, ModuleType): return name
    try: return getattr(torch.backends, name)
    except AttributeError: return None

def hasbackend(backend: str) -> bool:
    '''Check if a backend is both available and built.

    Parameters
    ----------
    backend : str
        The name of the backend to check.

    Returns
    -------
    bool
        True if the backend is available and built, False otherwise.
    '''
    be = getbackend(backend)
    return isbuilt(be) and isavail(be)

# %% ../nbs/60_bend.ipynb 26
def backend(name: str) -> Callable[[Callable], Callable[[], bool]]:
    '''Decorator to check if a specified backend is built and available.

    Parameters
    ----------
    name : str
        The name of the backend to check.

    Returns
    -------
    Callable[[Callable], bool]
        A decorator function that when applied to a function, returns the result of `hasbackend(name)`.
    '''
    def decorator(func: Callable) -> Callable[[], bool]:
        func.__doc__ = f'''Checks whether or not pytorch has {name} availble and was built with {name} in mind.'''
        @wraps(func, assigned=(__NAME__, __MODULE__, __QUALNAME__, __DOC__,))
        def wrapper(*args: P.args, **kwargs: P.kwargs) -> bool:
            return hasbackend(name)
        return wrapper
    return decorator

# %% ../nbs/60_bend.ipynb 28
@backend(MPS)
def hasmps() -> bool: 
    '''Checks whether or not pytorch has mps availble and was built with mps in mind.'''
    
@backend(CUDA)
def hascuda() -> bool:
    '''Checks whether or not pytorch has cuda availble and was built with cuda in mind.'''

# %% ../nbs/60_bend.ipynb 30
#| export
