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

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

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


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


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

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


# %% ../nbs/32_init.ipynb 15
try: import torch
except: ...

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


# %% ../nbs/32_init.ipynb 19
from nlit import DEVICE
from quac import intq, boolq, tensorq, deviceq
from chck import isnone, notnone
from asto import asdev
from etrc import RecurrentLayer

# %% ../nbs/32_init.ipynb 21
from .atyp import HiddenState, CellState, HiddenStates
from .dims import ndirected, rbatches_dim
from .tens import last_rn

# %% ../nbs/32_init.ipynb 24
def init_s0(
    s0: tensorq = None, 
    nlays: int = 1, 
    hsize: int = 1, 
    bsize: intq = None, 
    device: deviceq = None, 
    bidirectional: boolq = None
) -> tensorq:
    '''Initalize state tensor for a layer if `s0` is None.
    
    Parameters
    ----------
    s0: tensorq, 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 = ndirected(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/32_init.ipynb 26
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: tensorq, 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/32_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: tensorq, 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/32_init.ipynb 30
def init_r0(
    x: tensorq, 
    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 / cell state tensors for a recurrent layer if they are `None`.

    Parameters
    ----------
    x : tensorq
        The input tensor to the recurrent layer.

    h0 : tensorq, optional
        Initial hidden state, defaults to None.

    c0 : tensorq, optional
        Initial cell state (for LSTM), defaults to None.

    nlays : int, optional
        Number of layers in the RNN, defaults to 1.

    hcell : int, optional
        Number of cells in each layer, defaults to 1.

    hsize : intq, optional
        Hidden size, defaults to None.

    bsize : intq, optional
        Batch size, defaults to None.

    device : deviceq, optional
        Device to use, defaults to None.

    batch_first : bool, optional
        If True, batch dimension comes first, defaults to True.

    bidirectional : boolq, optional
        If True, RNN is bidirectional, defaults to None.
        
    kind : RecurrentLayer, optional
        Type of recurrent layer (LSTM, GRU, etc.), defaults to LSTM.

    Returns
    -------
    HiddenStates
        A tuple of the input tensor and initialized hidden and cell states.

    Examples
    --------
    >>> init_r0(torch.rand(3, 5, 10))
    # Returns the input tensor with initialized hidden and cell states.
    '''
    nlays = ndirected(nlays, bidirectional)
    
    bsize = bsize if notnone(bsize) else rbatches_dim(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/32_init.ipynb 32
def init_rn(x: tensorq, retlast: bool = True) -> tensorq:
    '''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.

    Parameters
    ----------
    x : tensorq
        The input tensor to the recurrent layer.
    retlast : bool, default: True
        If True, returns the last output of the tensor, otherwise returns the entire tensor, defaults to True.

    Returns
    -------
    tensorq
        The initialized output tensor of the recurrent layer.
        
    See Also
    --------
    last_rn : Get the last output of a recurrent layer.

    Examples
    --------
    >>> init_rn(torch.rand(3, 5, 10))
    # Returns either the last output or the entire tensor, depending on `retlast`.
    '''
    return last_rn(x) if retlast else x

# %% ../nbs/32_init.ipynb 34
def init_loss(x: tensorq = None, device: deviceq = None) -> float:
    dev = asdev(device or getattr(x, DEVICE, device))
    return torch.tensor(0.0).float().to(dev)

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