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

# %% auto 0
__all__ = ['is2d', 'is3d', 'as2d', 'as3d', 'as4d', 'last', 'batches', 'channels', 'get_torch_module']

# %% ../nbs/02_util.ipynb 6
from enum import Enum

# %% ../nbs/02_util.ipynb 8
from types import ModuleType

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

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


# %% ../nbs/02_util.ipynb 15
try: import torch, torch.nn as nn
except ImportError: torch, nn = None, None

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


# %% ../nbs/02_util.ipynb 19
from atyp import IntQ, BoolQ, Layer, Tensor, TensorQ, DeviceQ
from nchr import U1, NIL
from chck import isnone, notnone, isiter
from atup import argtup, jointups, nones

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


# %% ../nbs/02_util.ipynb 26
def is2d(x: Tensor) -> bool:
    '''Check if `t` is a 2D tensor.'''
    return x.ndim == 2

def is3d(x: Tensor) -> bool: 
    '''Check if `t` is a 3D tensor.'''
    return x.ndim == 3

# %% ../nbs/02_util.ipynb 28
def as2d(t: Tensor, device: DeviceQ = None) -> Tensor:
    return torch.flatten(t, start_dim=1).to(device)
    
def as3d(t: Tensor, device: DeviceQ = None) -> Tensor:
    return torch.flatten(t, start_dim=2).to(device)

def as4d(t: Tensor, height: int, width: int, channels: int = -1, device: DeviceQ = None) -> Tensor:
    shape = (height, width, )
    if is3d(t): t = torch.unflatten(t, 2, shape)
    else: t = torch.unflatten(t, 1, (channels, *shape, ))
    return t.to(device)

# %% ../nbs/02_util.ipynb 29
def last(t: Tensor, is_batched: bool = False) -> Tensor:
    '''Get the last value of the output `t`.'''
    if is2d(t) and not is_batched: return t[-1, :].unsqueeze(0)
    return t[:, -1, :].unsqueeze(1)

# %% ../nbs/02_util.ipynb 30
def batches(t: Tensor, batch_first: bool = True) -> IntQ: 
    '''Extract the batch size of tensor `t`.
    This will be either `None`, `x.size(0)` or `x.size(1)` depending on the 
    shape of `t` and the value of `batch_first`.

    Notes
    -----
    if `t` is a 2D tensor then `t` has shape:
        (L, H_in):    (seq_len, input_size)
    if `t` is a 3D tensor and `batch_first` is `True` then `t` has shape:
        (N, L, H_in): (batch_size, seq_len, input_size)
    otherwise `t` has shape:
        (L, N, H_in): (seq_len, batch_size, input_size)
    '''
    if is2d(t): return None
    return t.size(0 if batch_first else 1)

def channels(t: Tensor, batch_first: bool = True) -> IntQ:
    '''Extract the number of channels in tensor `t`.
    This will be either `x.size(0)` or `x.size(1)` depending on the 
    shape of `t` and the value of `batch_first`.
    
    Notes
    -----
    if `t` is a 2D tensor then `t` has shape:
        (L, H_in):    (seq_len, input_size)
    if `t` is a 3D tensor and `batch_first` is `True` then `t` has shape:
        (N, L, H_in): (batch_size, seq_len, input_size)
    otherwise `t` has shape:
        (L, N, H_in): (seq_len, batch_size, input_size)
    '''
    if is2d(t): return t.size(0)
    if is3d(t): return t.size(1 if batch_first else 0)
    return t.size(0)



# %% ../nbs/02_util.ipynb 33
def get_torch_module(
    obj, 
    *args,
    __module: ModuleType = nn, 
    __enumkey: BoolQ = None,
    __enumval: BoolQ = None, 
    __suffix: str = NIL,
    **kwargs
) -> Layer:  
    # 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)
        
    val = f'{val}{__suffix}'
    mod = getattr(__module, val)
    if args or kwargs: mod = mod(*args, **kwargs)
    return mod

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