# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/05_init.ipynb.

# %% auto 0
__all__ = ['directions', 'directed_layers', 'init_s0', 'init_h0', 'init_c0', 'init_r0', 'init_rn']

# %% ../nbs/05_init.ipynb 6
from enum import Enum

# %% ../nbs/05_init.ipynb 8
from types import ModuleType

# %% ../nbs/05_init.ipynb 11
#| export

# %% ../nbs/05_init.ipynb 13
#| export


# %% ../nbs/05_init.ipynb 15
try: import torch
except ImportError: torch = None

# %% ../nbs/05_init.ipynb 17
#| export


# %% ../nbs/05_init.ipynb 19
from atyp import IntQ, BoolQ, Layer, Tensor, TensorQ, DeviceQ
from chck import isnone, notnone
from nchr import U1, NIL

# %% ../nbs/05_init.ipynb 21
from .atyp import HiddenState, CellState, HiddenStates, HiddenStatesQ
from .util import last, batches
from .enum import RecurrentLayer

# %% ../nbs/05_init.ipynb 24
def directions(bidirectional: bool = False) -> int:
    '''Get the number of directions'''
    return int(2 if bidirectional else 1)

# %% ../nbs/05_init.ipynb 25
def directed_layers(num_layers: int = 1, bidirectional: bool = False) -> int:
    '''Calculate (D * L) where D is the number of directions and L is the number of layers.'''
    return num_layers * directions(bidirectional)

# %% ../nbs/05_init.ipynb 26
def init_s0(
    s0: TensorQ = None, nlays: int = 1, hsize: int = 1, bsize: IntQ = None, 
    device: DeviceQ = None, bidirectional: BoolQ = None
) -> Tensor:
    '''Initalize state tensor for a layer if `s0` is None.
    
    Parameters
    ----------
    s0: Tensor, default: None
        The initial state tensor. If None, the state tensor is initialized to zeros.
        
    nlays: int, default: 1
        The number of layers (D * L), where D is the number of directions (1 or 2) and L is the number of layers

    hsize: int, default: 1
        The number of hidden units in each layer

    bsize: int, default: None
        The number of bsize in the input data. If None, the input data is assumed to be 
        unbatched (i.e. a single sequence).

    device: torch.device, default: None
        The device to use for the tensor. If None, the default device is used.

    bidirectional: bool, default: None
        Whether or not to double the number of layers. If True, the number of directions is 2, otherwise 1.
    '''
    nlays = directed_layers(nlays, bidirectional)
    if notnone(s0): pass
    elif isnone(bsize): s0 = torch.zeros(nlays, hsize)
    else: s0 = torch.zeros(nlays, bsize, hsize)
    return s0.to(device or s0.device)
    

# %% ../nbs/05_init.ipynb 27
def init_h0(
    h0: TensorQ = None, nlays: int = 1, hsize: int = 1, bsize: IntQ = None, 
    device: DeviceQ = None, bidirectional: BoolQ = None
) -> HiddenState: 
    '''Initalize hidden state tensor for a layer if `h0` is None.
    
    Parameters
    ----------
    h0: Tensor, default: None
        The initial hidden state tensor. If None, the hidden state tensor is initialized to zeros.
        
    nlays: int, default: 1
        The number of layers (D * L), where D is the number of directions (1 or 2) and L is the number of layers

    hsize: int, default: 1
        The number of hidden units in each layer

    bsize: int, default: None
        The number of bsize in the input data. If None, the input data is assumed to be 
        unbatched (i.e. a single sequence).

    device: torch.device, default: None
        The device to use for the tensor. If None, the default device is used.
    
    bidirectional: bool, default: None
        Whether or not to double the number of layers. If True, the number of directions is 2, otherwise 1.
    '''
    return init_s0(h0, nlays, hsize, bsize, device, bidirectional)

# %% ../nbs/05_init.ipynb 28
def init_c0(
    c0: TensorQ = None, nlays: int = 1, hcell: int = 1, bsize: IntQ = None, 
    device: DeviceQ = None, bidirectional: BoolQ = None
) -> CellState: 
    '''Initalize cell state tensor for a layer if `c0` is None.
    
    Parameters
    ----------
    c0: Tensor, default: None
        The initial cell state tensor. If None, the cell state tensor is initialized to zeros.
        
    nlays: int, default: 1
        The number of layers (D * L), where D is the number of directions (1 or 2) and L is the number of layers

    hcell: int, default: 1
        The number of cell units in each layer

    bsize: int, default: None
        The number of bsize in the input data. If None, the input data is assumed to be 
        unbatched (i.e. a single sequence).

    device: torch.device, default: None
        The device to use for the tensor. If None, the default device is used.

    bidirectional: bool, default: None
        Whether or not to double the number of layers. If True, the number of directions is 2, otherwise 1.
    '''
    return init_s0(c0, nlays, hcell, bsize, device, bidirectional)

# %% ../nbs/05_init.ipynb 29
def init_r0(
    x: Tensor, h0: TensorQ = None, c0: TensorQ = None, 
    nlays: int = 1, hcell: int = 1, hsize: IntQ = None, bsize: IntQ = None, 
    device: DeviceQ = None, batch_first: bool = True, bidirectional: BoolQ = None, 
    kind: RecurrentLayer = RecurrentLayer.LSTM
) -> HiddenStates:
    '''Initialize the input tensor and the hidden and cell state tensors for a recurrent layer if they are None.'''
    nlays = directed_layers(nlays, bidirectional)
    
    bsize = bsize if notnone(bsize) else batches(x, batch_first)
    hsize = hsize if notnone(hsize) else hcell
    
    h0 = init_h0(h0, nlays, hsize, bsize, device, bidirectional=None)
    if kind == RecurrentLayer.LSTM:
        c0 = init_c0(c0, nlays, hcell, bsize, device, bidirectional=None)
        return x, (h0, c0)
    return x, h0

# %% ../nbs/05_init.ipynb 30
def init_rn(x: Tensor, retlast: bool = True) -> Tensor:
    '''Prepare the output tensor (i.e. rn) for a recurrent layer. If `retlast` is True, the last output is returned, otherwise the output is returned.'''
    return last(x) if retlast else x

# %% ../nbs/05_init.ipynb 32
#| export
