# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/02_util.ipynb.

# %% auto 0
__all__ = ['nextv', 'enumdict', 'evalue', 'missing', 'get_torch_module']

# %% ../nbs/02_util.ipynb 6
from enum import Enum, _EnumDict
from importlib import import_module

# %% ../nbs/02_util.ipynb 8
from types import ModuleType
from typing import (Type)

# %% ../nbs/02_util.ipynb 10
#| export


# %% ../nbs/02_util.ipynb 12
from nchr import NIL
from nlit import NAME, VALUE, _DEFAULT_
from quac import layer, boolq
from chck import isnone, notnone

# %% ../nbs/02_util.ipynb 14
from .cons import TORCH_NN
from .atyp import T

# %% ../nbs/02_util.ipynb 16
def nextv(name: str, start: int, count: int, last_values: list[str]):
    '''get the next value for an enum
    
    Parameters
    ----------
    name : str
        the name of the enum member
    start : int
        the starting value for the enum
        
    count : int
        the number of members in the enum
        
    last_values : list[str]
        the last values of the enum members
    
    Returns
    -------
    str
        the next value for the enum member
    
    See Also
    --------
    _EnumDict._generate_next_value : the default implementation
    '''
    return name


def enumdict(name: str) -> _EnumDict:
    '''create a new enum dict
    
    Parameters
    ----------
    name : str
        the name of the enum class
        
    Returns
    -------
    _EnumDict
        the new enum dict
        
    Notes
    -----
    This function is a workaround for the fact that the `Enum` class does not expose the `_EnumDict` class.
    Thus the paramter `name` is used to set the `_cls_name` attribute of the new enum dict, which is required.
    
    '''
    newdict = _EnumDict()
    newdict._cls_name = name
    newdict._generate_next_value = nextv
    # try: newdict._generate_next_value = StrEnum._generate_next_value_
    # except: newdict._generate_next_value = Enum._generate_next_value_
    return newdict

def evalue(e: Enum): 
    '''Get the value of an enum member or return the enum member itself.
    This is useful if you are unsure if `e` is the enum or an enum member.
    '''
    return getattr(e, VALUE, e)

# %% ../nbs/02_util.ipynb 17
def missing(enm: Type[T], val: str | T, default = None) -> T:
    '''Method called for missing members during lookup.

    Parameters
    ----------
    cls : Enum
        The enumeration class.
    
    val : str
        The value to check for in the enum members.
    

    Returns
    -------
    Enum member or default
        The matched enum member or the default value.
    '''
    # for mem in enm:
    #     mal: member = getattr(mem, VALUE, '')
    #     if str(mal).casefold() == str(val).casefold(): return mem
    # if hasattr(enm, '_default_'): return enm._default_() or default
    # return default

    for mem in enm:
        mal = getattr(mem, VALUE, NIL)        
        if str(mal).casefold() == str(val).casefold(): return mem
    if hasattr(enm, _DEFAULT_): 
        edef = enm._default_()
        try: return enm(edef)
        except: ...
        try: return enm(default)
        except: return default
    return default

# %% ../nbs/02_util.ipynb 18
def get_torch_module(
    obj, *args, __module: ModuleType = TORCH_NN, 
    __enumkey: boolq = None, __enumval: boolq = None, __suffix: str = NIL,
    **kwargs
) -> layer:
    '''
    Fetch a torch module based on an enumeration value or its name.

    Parameters
    ----------
    obj : Enum
        An enumeration instance or class.
    args : tuple
        Positional arguments to pass to the module constructor.
    __module : ModuleType, optional
        The torch module to fetch from, default is `torch.nn`.
        
    __enumkey : BoolQ, optional
        If True, use the enum's name for fetching, otherwise use the value.
        
    __enumval : BoolQ, optional
        If True, use the enum's value for fetching, otherwise use the name.
        
    __suffix : str, optional
        Suffix to append to the enum value or name.
        
    kwargs : dict
        Keyword arguments to pass to the module constructor.

    Returns
    -------
    Layer
        The instantiated torch module.
    '''
    # ensure that only one of __usekey and __useval is True
    usekey, useval = None, None
    if notnone(__enumkey): usekey = __enumkey
    if notnone(__enumval): useval = __enumval
    if isnone(usekey) and isnone(useval): usekey = True
    
    usekey = not useval if notnone(useval) else usekey
    useval = not usekey if notnone(usekey) else useval
    
    if issubclass(type(obj), Enum):
        val = getattr(obj, NAME if usekey else VALUE, obj.value)
    else: val = obj
        
    val = f'{val}{__suffix}'
    
    if __module is None: __module = TORCH_NN
    if isinstance(__module, str): __module = import_module(__module)
    mod = getattr(__module, val)
    if args or kwargs: mod = mod(*args, **kwargs)
    return mod

# %% ../nbs/02_util.ipynb 20
#| export
