# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/03_enum.ipynb.

# %% auto 0
__all__ = ['missing', 'Channels', 'NonLinearity', 'InitMethod', 'RecurrentLayer']

# %% ../nbs/03_enum.ipynb 6
from enum import Enum, StrEnum, auto

# %% ../nbs/03_enum.ipynb 8
#| export


# %% ../nbs/03_enum.ipynb 11
#| export

# %% ../nbs/03_enum.ipynb 13
try:
    import numpy as np
except ImportError:
    ...

# %% ../nbs/03_enum.ipynb 15
try:
    import torch, torch.nn as nn
except ImportError:
    ...

# %% ../nbs/03_enum.ipynb 17
#| export


# %% ../nbs/03_enum.ipynb 19
from nchr import U1, NIL
from atyp import Tensor, Layer, DeviceQ, DTypeQ

# %% ../nbs/03_enum.ipynb 21
from .util import get_torch_module

# %% ../nbs/03_enum.ipynb 24
def missing(cls, val: str, default = None):
    val = val.lower()
    for member in cls:
        if member.value.lower() == val: 
            return member
    return default

# %% ../nbs/03_enum.ipynb 26
class Channels(StrEnum):
    THWC = auto() # 0, 1, 2, 3
    TCHW = auto() # 0, 3, 1, 2
    
    @classmethod
    def flip(cls, t: Tensor) -> Tensor: 
        return t.transpose(1, 3)
    
    @classmethod
    def cidx(cls, t: Tensor) -> int:
        '''Returns the index of the channel dimension'''
        idx = -1
        for i in range(len(t.shape)):
            if t.shape[i] != 3: continue
            idx = i
        
        idx = idx if idx >= 0 else np.argmin(t.shape)
        if (val := t.shape[idx]) not in {1, 3}: 
            raise ValueError(f'Expected eitehr 1 or 3 channels, got {val}')
        return idx
    
    @classmethod
    def _missing_(cls, val: str): return missing(cls, val, cls.THWC)
    
    @classmethod
    def safe(cls, c: 'Channels'):
        try: return cls(c)
        except ValueError: return cls.THWC
        
    @classmethod
    def to(cls, t: Tensor, c: 'Channels' = 'THWC') -> Tensor:
        c = cls.safe(c)
        cidx = cls.cidx(t)
        if c == cls.THWC and cidx == 1: return cls.flip(t)
        if c == cls.THWC and cidx == 3: return t
        if c == cls.TCHW and cidx == 1: return t
        if c == cls.TCHW and cidx == 3: return cls.flip(t)
        return t

# %% ../nbs/03_enum.ipynb 28
class NonLinearity(StrEnum):
    # weighted sum, nonlinearity
    ELU = auto()
    Hardshrink = auto()
    Hardsigmoid = auto()
    Hardtanh = auto()
    Hardswish = auto()
    LeakyReLU = auto()
    LogSigmoid = auto()
    MultiheadAttention = auto()
    PReLU = auto()
    SELU = auto()
    CELU = auto()
    GELU = auto()
    Sigmoid = auto()
    SiLU = auto()
    Mish = auto()
    Softplus = auto()
    Softshrink = auto()
    Tanh = auto()
    Tanhshrink = auto()
    Threshold = auto()
    GLU = auto()
    
    # other
    Softmin = auto()
    Softmax = auto()
    Softmax2d = auto()
    LogSoftmax = auto()
    AdaptiveLogSoftmaxWithLoss = auto()

    @classmethod
    def _missing_(cls, val: str): return missing(cls, val, cls.Tanh)
    
    def get(self, *args, **kwargs) -> Layer:
        return get_torch_module(self, *args, **kwargs)

# %% ../nbs/03_enum.ipynb 30
class InitMethod(StrEnum):
    constant = auto()
    dirac = auto()
    eye = auto()
    kaiming_normal = auto()
    normal = auto()
    orthogonal = auto()
    sparse = auto()
    uniform = auto()
    xavier_normal = auto()
    xavier_uniform = auto()
    zeros = auto()

    @classmethod
    def _missing_(cls, val: str): return missing(cls, val, cls.xavier_normal)

    def get(self, *args, **kwargs) -> Layer:
        suffix = kwargs.pop('__suffix', U1)
        return get_torch_module(self, *args, __module=nn.init, __enumval=True, __suffix=suffix, **kwargs)

# %% ../nbs/03_enum.ipynb 32
class RecurrentLayer(StrEnum):
    GRU = auto()
    RNN = auto()
    LSTM = auto()
    @classmethod
    def _missing_(cls, val: str): return missing(cls, val, cls.LSTM)
    def get(self, *args, **kwargs) -> Layer:        
        return get_torch_module(self, *args, **kwargs)

# %% ../nbs/03_enum.ipynb 34
#| export
