# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/31_tens.ipynb.

# %% auto 0
__all__ = ['torchnans', 'nanslike', 'nonzerosfrom', 'unpadded_len', 'pack_padded', 'sort_sequences', 'pad', 'pack', 'unpack',
           'augzeros', 'augment_zeros', 'last', 'last_rn']

# %% ../nbs/31_tens.ipynb 6
from functools import wraps, reduce

# %% ../nbs/31_tens.ipynb 8
from types import ModuleType

# %% ../nbs/31_tens.ipynb 11
#| export

# %% ../nbs/31_tens.ipynb 13
try: import numpy as np
except ImportError: ...

try: import pandas as pd
except ImportError: ...

# %% ../nbs/31_tens.ipynb 15
try: import torch
except: ...

try: 
    from torch.nn.utils.rnn import (
        pad_sequence, pack_padded_sequence, 
        pad_packed_sequence, PackedSequence
    )
except: ...

# %% ../nbs/31_tens.ipynb 17
#| export


# %% ../nbs/31_tens.ipynb 19
from quac import (tensor, inttensor, inttensorq)
from chck import isnone
from asto import nanslike as _nanslike, torchnans as _torchnans

# %% ../nbs/31_tens.ipynb 21
from .cons import WRAPS_ASSIGN_ANNDOCS
from .dims import adjsize, is2d

# %% ../nbs/31_tens.ipynb 24
@wraps(_torchnans, assigned=WRAPS_ASSIGN_ANNDOCS)
def torchnans(t: tensor) -> tensor:
    return _torchnans(t)

@wraps(_nanslike, assigned=WRAPS_ASSIGN_ANNDOCS)
def nanslike(t: tensor) -> tensor:
    return _nanslike(t)

# %% ../nbs/31_tens.ipynb 27
def nonzerosfrom(t: tensor, axis: int = 1) -> inttensor:
    '''
    Count the number of non-zero elements in each sequence along a specified axis.

    Parameters
    ----------
    t : tensor
        The input tensor containing sequences.
    axis : int, optional
        The axis along which to count non-zero elements, default is 1.

    Returns
    -------
    inttensor
        An integer tensor indicating the count of non-zero elements in each sequence.

    Examples
    --------
    >>> nonzerosfrom(torch.tensor([[1, 2, 0], [3, 0, 0]]))
    tensor([2, 1])
    '''
    return reduce(lambda t, _: t.count_nonzero(-1), t.shape[axis:], t)

# %% ../nbs/31_tens.ipynb 29
def unpadded_len(t: tensor, lens: inttensorq = None) -> inttensor:
    '''
    Determine the unpadded length of each sequence in a batch of sequences.

    Parameters
    ----------
    t : tensor
        The tensor containing sequences, possibly with padding.
    lens : inttensorq, optional
        Precomputed lengths of the sequences, defaults to None.

    Returns
    -------
    inttensor
        An integer tensor representing the length of each sequence.

    Examples
    --------
    >>> unpadded_len(torch.tensor([[1, 2, 0], [3, 0, 0]]))
    tensor([2, 1])
    '''
    return nonzerosfrom(t, axis=1) if isnone(lens) else lens

# %% ../nbs/31_tens.ipynb 31
def pack_padded(p: tensor, lens: inttensorq = None) -> PackedSequence:
    '''
    Convert a batch of padded sequences to a PackedSequence.

    Parameters
    ----------
    p : tensor
        The batch of padded sequences.
    lens : inttensorq, optional
        The lengths of each sequence in the batch, defaults to None.

    Returns
    -------
    PackedSequence
        The packed sequence.

    Examples
    --------
    >>> p = torch.tensor([[1, 2, 0], [3, 4, 0]])
    >>> pack_padded(p)
    PackedSequence(...)
    '''
    lens = unpadded_len(p, lens).clone().cpu()
    return pack_padded_sequence(p, lens, batch_first=True)

# %% ../nbs/31_tens.ipynb 33
def sort_sequences(t: tensor, lens: inttensorq = None) -> tuple[tensor, inttensor, inttensor]:
    '''
    Sort a batch of sequences by their lengths in descending order.

    Parameters
    ----------
    t : tensor
        The batch of sequences.
    lens : inttensorq, optional
        The lengths of each sequence in the batch, defaults to None.

    Returns
    -------
    tensor
        The sorted sequences.
    inttensor
        The lengths of the sorted sequences.
    inttensor
        The indices that sort the batch.

    Examples
    --------
    >>> t = torch.tensor([[1, 2, 0], [3, 4, 5]])
    >>> sorted_t, sorted_lens, sorted_idx = sort_sequences(t)
    '''
    lens = unpadded_len(t, lens)
    sorted_lens, sorted_idx = lens.sort(descending=True)
    sorted_t = t[sorted_idx]
    return sorted_t, sorted_lens, sorted_idx

# %% ../nbs/31_tens.ipynb 35
def pad(t: tensor, batch_first: bool = True, value: float = 0) -> tensor:
    '''
    Pad a batch of variable-length sequences to the same length with a specified value.

    This function uses `torch.nn.utils.rnn.pad_sequence` to pad a batch of sequences. 
    It's useful when dealing with batches of sequences of varying lengths, 
    particularly in preparation for neural network inputs like RNNs.

    Parameters
    ----------
    t : tensor
        A tensor or a list of tensors representing sequences of varying lengths.
    batch_first : bool, optional
        If True, the output tensor will have the batch dimension as the first dimension, 
        otherwise the sequence length dimension will be first. Defaults to True.
    value : float, optional
        The padding value to use to fill the sequences to match the longest sequence. Defaults to 0.

    Returns
    -------
    tensor
        The padded tensor with sequences of the same length.

    Examples
    --------
    >>> seqs = [torch.tensor([1, 2]), torch.tensor([3, 4, 5])]
    >>> pad(seqs)
    tensor([[1, 2, 0],
            [3, 4, 5]])
    '''
    return pad_sequence(t, batch_first=batch_first, padding_value=value)

# %% ../nbs/31_tens.ipynb 37
def pack(t: tensor, lens: inttensorq = None) -> tuple[PackedSequence, inttensor]:
    '''
    Packs a batch of sequences, sorting them by length in descending order.

    Parameters
    ----------
    t : tensor
        The batch of sequences to pack.
    lens : inttensorq, optional
        The lengths of each sequence in the batch, defaults to None.

    Returns
    -------
    PackedSequence
        The packed sequence.
    inttensor
        The indices that sort the batch.

    Examples
    --------
    >>> t = torch.tensor([[1, 2, 0], [3, 4, 5]])
    >>> packed, sorted_idx = pack(t)
    '''
    sorted_x, sorted_lens, sorted_idx = sort_sequences(t, lens)
    packed = pack_padded(sorted_x, sorted_lens)
    return packed, sorted_idx

# %% ../nbs/31_tens.ipynb 39
def unpack(p: PackedSequence) -> tuple[tensor, inttensor]:
    '''
    Converts a PackedSequence back to a padded sequence.

    Parameters
    ----------
    p : PackedSequence
        The packed sequence.

    Returns
    -------
    tensor
        The padded sequence tensor.
    inttensor
        The lengths of each sequence in the padded tensor.

    Examples
    --------
    >>> p = pack(torch.tensor([[1, 2, 0], [3, 4, 5]]))[0]
    >>> padded, lens = unpack(p)
    '''
    padded, lens = pad_packed_sequence(p, batch_first=True)
    return padded, lens  

# %% ../nbs/31_tens.ipynb 41
def augzeros(x: tensor, idx: int = 1, dim: int = 1, add: bool = False) -> tensor:
    '''Creates a zero tensor matching the adjusted size of tensor `x` with a new dimension `dim` at index `idx`.
    
    Parameters
    ----------
    x : tensor
        The reference tensor for size adjustment.
        
    idx : int, default 1
        The index at which to adjust the size.
        
    dim : int, default 1
        The dimension to adjust at `idx`.
        
    add : bool, default False
        Whether to add `dim` as a new dimension at `idx` or to replace the existing dimension.
        
    Returns
    -------
    tensor
        A zero tensor with the adjusted size.
        
    See Also
    --------
    adjsize : Adjust the size of a tensor.
    augzeros : Alias for augzeros
    '''
    try: import torch
    except: ...
    dims = adjsize(x, idx, dim, add)
    return torch.zeros(dims)


@wraps(augzeros)
def augment_zeros(*args, **kwargs):
    '''Alias for augzeros'''
    return augzeros(*args, **kwargs)

# %% ../nbs/31_tens.ipynb 43
def last(t: tensor, batch_idx: int = 0, is_batched: bool = False, as_batched: bool = False) -> tensor:
    '''Extract the last element from a specified batch dimension of a PyTorch tensor.

    Parameters
    ----------
    t : torch.Tensor
        The tensor from which to extract the last element.
        
    batch_idx : int, optional
        The index of the dimension representing the batch.
        
    is_batched : bool, optional
        A flag indicating whether the tensor is batched.
        
    as_batched : bool, optional
        A flag indicating whether to return the result as a batched tensor.

    Returns
    -------
    torch.Tensor
        The last element of the specified batch dimension of the tensor.
    '''
    if not is_batched: batch_idx = 0        
    # If the tensor is batched, extract the last element from the specified batch dimension
    slcs = [slice(None)] * len(t.shape)
    # Set the slicing object to take the last index in the batch dimension
    slcs[batch_idx] = -1
    res = t[tuple(slcs)]
    return res.unsqueeze(batch_idx) if as_batched else res

# %% ../nbs/31_tens.ipynb 44
def last_rn(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/31_tens.ipynb 46
#| export
