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

# %% auto 0
__all__ = ['T', 'O', 'P', 'Ts', 'VidsInfo', 'XYArray', 'LabelArray', 'DistanceArray', 'RotationMatrix', 'DynamicsInput',
           'DynamicsOutput', 'CellState', 'HiddenState', 'RNNState', 'GRUState', 'LSTMStates', 'RNNStateQ', 'GRUStateQ',
           'LSTMStatesQ', 'RecurrentStates', 'RecurrentStatesQ', 'HiddenStates', 'HiddenStatesQ', 'TorchDynForward',
           'DynamicsForward', 'KineticsForward', 'LossFunction', 'DiffusionFunction', 'star']

# %% ../nbs/02_atyp.ipynb 6
#| export

# %% ../nbs/02_atyp.ipynb 8
from typing import (
    Self, Union, Iterable, 
    TypeVar, ParamSpec, TypeAlias, TypeVarTuple, 
    Optional, Callable, ForwardRef
)

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


# %% ../nbs/02_atyp.ipynb 12
from nchr import STAR
from quac import (
    ints, dim1, dim2, num, bbox, lit, size, tensor, path,
    neuraldynamics
)

# %% ../nbs/02_atyp.ipynb 14
#| export


# %% ../nbs/02_atyp.ipynb 16
T = TypeVar('T')
'''Type variable for generic types''';

O = TypeVar('O')
'''Type variable for generic types reprenesting "other"''';

P = ParamSpec('P')
'''Parameter specification type variable, capturing the type of function parameters.''';

Ts = TypeVarTuple('Ts')

# %% ../nbs/02_atyp.ipynb 19
_VidInfo: TypeAlias = ForwardRef('VidInfo', module='utrc.kwds', is_class=True)

VidsInfo: TypeAlias = dict[path, _VidInfo]
'''Dictionary of (path, vidinfo) pairs''';

# %% ../nbs/02_atyp.ipynb 21
class star(lit, __lit__ = STAR): 
    ...

XYArray: TypeAlias = dim2[float, float]

LabelArray: TypeAlias = dim1[num]

DistanceArray = dim1[float]

RotationMatrix = bbox[size[2, 2], float]

# %% ../nbs/02_atyp.ipynb 24
_X: TypeAlias = tensor
'''`x`''';

_T: TypeAlias = tensor
'''`t`''';

_Save_At: TypeAlias = Iterable
'''`save_at`''';

_Args: TypeAlias = dict
'''`args`''';

_AugmentLayer: TypeAlias = ForwardRef('AugmentLayer', module='ktrc.lays', is_class=True)
'''`aug`''';

_Dynamics: TypeAlias = neuraldynamics
'''`dyn`''';

_DropAugLayer: TypeAlias = ForwardRef('DropAugLayer', module='ktrc.lays', is_class=True)
'''`dec`''';

_Deltas: TypeAlias = ints
'''`deltas`''';

_UseAug: TypeAlias = bool
'''`useaug`''';

_UseDec: TypeAlias = bool
'''`usedec`''';

_Y: TypeAlias = tensor
'''`y`''';

_TY: TypeAlias = tuple[_T, _Y]
'''`(t, y)`''';

# %% ../nbs/02_atyp.ipynb 26
DynamicsInput: TypeAlias = tuple[_X, Optional[_T], _Save_At, _Args]
'''`x, t, save_at, args`''';

DynamicsOutput: TypeAlias = Union[_Y, tuple[_T, _Y]]
'''`y` or `(t, y)`''';

# %% ../nbs/02_atyp.ipynb 28
CellState: TypeAlias = tensor
'''Type alias for representing the cell state in LSTM layers (`torch.Tensor`).

Represents the cell state tensor in LSTM layers, which carries information 
across cell sequences during the operation of the LSTM.
''';

HiddenState: TypeAlias = tensor
'''Type alias for representing the hidden state in RNN layers (`torch.Tensor`).

Represents the hidden state tensor in RNN layers, including LSTM and GRU, 
which is passed across different time steps of the RNN operation.
''';

RNNState: TypeAlias = HiddenState
'''Type alias for the hidden state in a simple RNN layer (`torch.Tensor`).

Represents the hidden state in a simple RNN (Recurrent Neural Network) layer, 
holding the information passed across time steps.
''';

GRUState: TypeAlias = HiddenState
'''Type alias for the hidden state in a GRU (Gated Recurrent Unit) layer  (`torch.Tensor`).

Represents the hidden state in a GRU layer, which is a more complex RNN 
architecture that includes gating mechanisms.
''';

LSTMStates: TypeAlias = tuple[HiddenState, CellState]
'''Type alias for the combined hidden and cell states in an LSTM layer (`tuple[torch.Tensor, torch.Tensor]`).

Represents a tuple containing both the hidden state and cell state in an LSTM 
(Long Short-Term Memory) layer, which are used for carrying information across 
time steps and sequences.
''';

RNNStateQ: TypeAlias = Optional[RNNState]
'''Type alias for an optional hidden state in a simple RNN layer (`Optional[torch.Tensor]`).

Represents an optional hidden state for a simple RNN layer, which may or may not 
be provided depending on the context of the RNN operation.
''';

GRUStateQ: TypeAlias = Optional[GRUState]
'''Type alias for an optional hidden state in a GRU layer (`Optional[torch.Tensor]`).

Represents an optional hidden state for a GRU (Gated Recurrent Unit) layer, 
which may be omitted in certain use cases or initial states.
''';

LSTMStatesQ: TypeAlias = Optional[LSTMStates]
'''Type alias for optional combined hidden and cell states in an LSTM layer
(`Optional[tuple[torch.Tensor, torch.Tensor]]`).

Represents an optional tuple of hidden and cell states for an LSTM (Long Short-Term Memory) 
layer, which can be omitted in initial steps or specific applications.
''';

RecurrentStates: TypeAlias = Union[RNNState, GRUState, LSTMStates]
'''Type alias for representing the hidden states of various recurrent layers
`torch.Tensor | tuple[torch.Tensor, torch.Tensor`.

Encapsulates the hidden states for different types of recurrent layers 
(simple RNN, GRU, and LSTM), providing a unified type for handling 
recurrent layer states.
'''

RecurrentStatesQ: TypeAlias = Optional[RecurrentStates]
'''Type alias for an optional representation of hidden states in recurrent layers
`Optional[torch.Tensor | tuple[torch.Tensor, torch.Tensor]`.

Represents an optional type that can encapsulate hidden states of various 
recurrent layers, allowing for flexibility in specifying these states.
'''

HiddenStates: TypeAlias = RecurrentStates
'''Type alias for hidden states in any recurrent neural network layers
`torch.Tensor | tuple[torch.Tensor, torch.Tensor`.

A general type for representing hidden states in RNN, GRU, or LSTM layers,
facilitating a consistent interface for these various architectures.
'''

HiddenStatesQ: TypeAlias = RecurrentStatesQ
'''Type alias for an optional hidden state in recurrent neural network layers
`Optional[torch.Tensor | tuple[torch.Tensor, torch.Tensor]`.

Provides an optional type for representing hidden states in RNN, GRU, or LSTM 
layers, which can be used in contexts where the state may not be initially defined.
''';

# %% ../nbs/02_atyp.ipynb 31
TorchDynForward: TypeAlias = Callable[[DynamicsInput], DynamicsOutput]
'''`x: tensor, t: Optional[tensor] = None, save_at: Iterable = (), args: dict = {}` -> `y` or `(t, y)`''';

DynamicsForward: TypeAlias = Callable[
    [
        _X, Optional[_T], _Save_At, _Args,
        # Unpack[DynamicsInput],
        Optional[_AugmentLayer],
        Optional[_Dynamics],
        Optional[_DropAugLayer],
        Optional[_Deltas], 
        _UseAug, 
        _UseDec
    ], 
    DynamicsOutput
]
'''`forward(\
    x: tensor,\
    t: Optional[tensor] = None,\
    save_at: Iterable = (),\
    args: dict = {},\
    aug: Optional[AugmentLayer] = None,\
    dyn: Optional[neuraldynamics] = None,\
    dec: Optional[DropAugLayer] = None,\
    deltas: Optional[list[int]] = None,\
    useaug: bool = True,\
    usedec: bool = True\
)` -> `y` or `(t, y)`
''';

KineticsForward: TypeAlias = Callable[
    [
        Self, _X, Optional[_T], _Save_At, _Args,
        Optional[_Deltas], _UseAug, _UseDec
    ], 
    DynamicsOutput
]
'''`dynamics(self,\
    x: tensor,\
    t: Optional[tensor] = None,\
    save_at: Iterable = (),\
    args: dict = {},\
    deltas: Optional[list[int]] = None,\
    useaug: bool = True,\
    usedec: bool = True\
)` -> `y` or `(t, y)`
''';

# %% ../nbs/02_atyp.ipynb 33
LossFunction: TypeAlias = Callable[[tensor, tensor], Union[tensor,  num]]

# %% ../nbs/02_atyp.ipynb 35
DiffusionFunction: TypeAlias = Union[
    'Diffusion', 'DiffusionMap', 'DiffusionDistance', 
    'DiffusionAffinity', 'PHATEDistance'
]

# %% ../nbs/02_atyp.ipynb 37
#| export
