# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/30_dims.ipynb.

# %% auto 0
__all__ = ['is2d', 'is3d', 'as2d', 'as3d', 'as4d', 'iasdim', 'idx_as_dim', 'augidx', 'augment_idx', 'adjsize', 'adjust_size',
           'insert_dim', 'fcs', 'fce', 'frs', 'fre', 'getbbox', 'vidpxs', 'subpxs', 'rbatches_dim', 'rchannels_dim',
           'ndirections', 'ndirected']

# %% ../nbs/30_dims.ipynb 6
from functools import wraps

# %% ../nbs/30_dims.ipynb 8
#| export


# %% ../nbs/30_dims.ipynb 11
#| export

# %% ../nbs/30_dims.ipynb 13
#| export


# %% ../nbs/30_dims.ipynb 15
try: import torch
except: ...

# %% ../nbs/30_dims.ipynb 17
#| export


# %% ../nbs/30_dims.ipynb 19
from quac import tensor, deviceq, intq
from quac import (vidcap, pathlike, bbox as BBox, rect)
from chck import isint, isnone, ispathlike
from etrc import SliceFormat

# %% ../nbs/30_dims.ipynb 21
from .cons import WRAPS_ASSIGN_ANNDOCS
from .attr import _shape, _ndim

# %% ../nbs/30_dims.ipynb 24
def is2d(t: tensor) -> bool:
    '''Check if `t` is a 2D tensor.'''
    return _ndim(t) == 2

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

# %% ../nbs/30_dims.ipynb 26
def as2d(t: tensor, device: deviceq = None) -> tensor:
    '''Flatten a tensor to 2D.'''
    try: import torch
    except: ...
    return torch.flatten(t, start_dim=1).to(device)
    
def as3d(t: tensor, device: deviceq = None) -> tensor:
    '''Flatten a tensor to 3D.'''
    try: import torch
    except: ...
    return torch.flatten(t, start_dim=2).to(device)

def as4d(t: tensor, height: int, width: int, channels: int = -1, device: deviceq = None) -> tensor:
    '''Unflatten a tensor to 4D.'''
    try: import torch
    except: ...
    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/30_dims.ipynb 28
def iindims(i: int, dims: tuple[int, ...]) -> bool:
    '''Ensure that the index `i` is within the bounds of the `dims`.
    
    Parameters
    ----------
    i : int
        The index to check.
    dims : tuple[int, ...]
        The dimensions tuple to check against.
        
    Returns
    -------
    bool
        True if `i` is within the bounds of `dims`, False otherwise.
        
    See Also
    --------
    is_idx_in_dims : Alias for `iindims`.
    '''
    return abs(i) < len(dims)

@wraps(iindims, assigned=WRAPS_ASSIGN_ANNDOCS)
def is_idx_in_dims(i: int, dims: tuple[int, ...]) -> bool:
    '''Alias for `iindims`.'''
    return iindims(i, dims)

# %% ../nbs/30_dims.ipynb 30
def iasdim(i: int, dims: tuple[int, ...]) -> int:
    '''Ensure that the index `i` is within the bounds of the `dims` or return -1.
    
    Parameters
    ----------
    i : int
        The index to adjust.
        
    dims : tuple[int, ...]
        The dimensions tuple to check against.
        
    Returns
    -------
    int
        The adjusted index if within bounds, otherwise -1.
        
    See Also
    --------
    iindims : Check if an index is within the bounds of the dimensions.
    idx_as_dim : Alias for `iasdim`.
    '''
    return i if iindims(i, dims) else -1

@wraps(iasdim, assigned=WRAPS_ASSIGN_ANNDOCS)
def idx_as_dim(i: int, dims: tuple[int, ...]) -> bool:
    '''Alias for `iasdim`.'''
    return iasdim(i, dims)

# %% ../nbs/30_dims.ipynb 32
def augidx(i: int, dims: int | tuple[int, ...]) -> int:
    '''Safely determine the augment index.'''
    ndim = dims if isint(dims) else len(dims) 
    aidx = aidx if abs(aidx) < ndim else -1
    return aidx

@wraps(augidx, assigned=WRAPS_ASSIGN_ANNDOCS)
def augment_idx(i: int, dims: tuple[int, ...]) -> bool:
    '''Alias for `augidx`.'''
    return augidx(i, dims)

# %% ../nbs/30_dims.ipynb 34
def adjsize(x: tensor, idx: int = 1, dim: int = 1, add: bool = False) -> tuple[int, ...]:
    '''Adjusts the size of tensor `x` by setting or adding a dimension `dim` at index `idx`.
    
    Return the size of the tensor `x` if adjusted at index `idx` to `dim`.
    
    Parameters
    ----------
    x : tensor
        The tensor to adjust the size of.
        
    idx : int, default: 1
        The index at which to set or add the dimension.
        
    dim : int, default; 1
        The dimension to set or add.
        
    add : bool, default: False
        If True, adds `dim` as a new dimension at `idx`; if False, sets `dim` at `idx`.
        
    Returns
    -------
    tuple
        The adjusted dimensions as a tuple.
        
    See Also
    --------
    insert_dim : Alias for `adjsize`.
    adjust_size : Alias for `adjsize`.
    '''
    dims = list(_shape(x))
    aidx = augidx(idx, dims)
    dims =  dims + [None]
    
    # NOTE: dim will never overwrite a dimension if add is True
    if add: dims = dims[:aidx] + [dim,] + dims[aidx:]
        
    # NOTE: dim might overwrite a dimension if add is False
    else: dims[aidx] = dim
    
    if isnone(dims[-1]): dims.pop()
        
    return tuple(dims)

@wraps(adjsize, assigned=WRAPS_ASSIGN_ANNDOCS)
def adjust_size(x: tensor, idx: int = 1, dim: int = 1, add: bool = False) -> tuple[int, ...]:
    '''Alias for `adjsize`.'''
    return adjsize(x, idx, dim, add)

@wraps(adjsize, assigned=WRAPS_ASSIGN_ANNDOCS)
def insert_dim(x: tensor, idx: int = 1, dim: int = 1, add: bool = False) -> tuple[int, ...]:
    '''Alias for `adjsize`.'''
    return adjsize(x, idx, dim, add)

# %% ../nbs/30_dims.ipynb 37
def fcs(col: int, width: int)  -> int: 
    '''
    Calculate the starting pixel column for a frame in a grid layout.

    Parameters
    ----------
    col : int
        Column index of the frame in the grid.
    width : int
        Width of each frame in the grid.

    Returns
    -------
    int
        Starting pixel column for the specified frame.

    Examples
    --------
    >>> fcs(1, 100)  # Start column for the second frame in a grid with 100px wide frames
    100
    '''
    return (col + 0) * width


# %% ../nbs/30_dims.ipynb 39
def fce(col: int, width: int)  -> int: 
    '''
    Calculate the ending pixel column for a frame in a grid layout.

    Parameters
    ----------
    col : int
        Column index of the frame in the grid.
    width : int
        Width of each frame in the grid.

    Returns
    -------
    int
        Ending pixel column for the specified frame.

    Examples
    --------
    >>> fce(1, 100)  # End column for the second frame in a grid with 100px wide frames
    200
    '''
    return (col + 1) * width

# %% ../nbs/30_dims.ipynb 41
def frs(row: int, height: int) -> int: 
    '''
    Calculate the starting pixel row for a frame in a grid layout.

    Parameters
    ----------
    row : int
        Row index of the frame in the grid.
    height : int
        Height of each frame in the grid.

    Returns
    -------
    int
        Starting pixel row for the specified frame.

    Examples
    --------
    >>> frs(1, 50)  # Start row for the second frame in a grid with 50px high frames
    50
    '''
    return (row + 0) * height

# %% ../nbs/30_dims.ipynb 43
def fre(row: int, height: int) -> int: 
    '''
    Calculate the ending pixel row for a frame in a grid layout.

    Parameters
    ----------
    row : int
        Row index of the frame in the grid.
    height : int
        Height of each frame in the grid.

    Returns
    -------
    int
        Ending pixel row for the specified frame.

    Examples
    --------
    >>> fre(1, 50)  # End row for the second frame in a grid with 50px high frames
    100
    '''
    return (row + 1) * height

# %% ../nbs/30_dims.ipynb 45
def getbbox(row: int, col: int, width: int, height: int, fmt: SliceFormat = SliceFormat.coord) -> BBox:
    '''
    Calculate the bounding box for a subframe in a grid layout (defaults in position format)

    Parameters
    ----------
    row : int
        Row index of the subframe in the grid.

    col : int
        Column index of the subframe in the grid.

    width : int
        Width of each frame in the grid.

    height : int
        Height of each frame in the grid.

    fmt : SliceFormat, optional
        Format for the bounding box, defaults to `SliceFormat.coord`.

    Returns
    -------
    BBox
        Bounding box for the specified subframe.

    Examples
    --------
    >>> getbbox(2, 3, 200, 400)
    ((800, 600), (1200, 800)) # ((x0, y0), (x1, y1))

    >>> getbbox(1, 1, 100, 50, SliceFormat.slice)
    ((800, 1200), (600, 800)) # ((x0, x1), (y0, y1))
    '''
    bbox = (frs(row, height), fcs(col, width)), (fre(row, height), fce(col, width))
    return SliceFormat.flip(*bbox) if fmt == SliceFormat.slice else bbox

# %% ../nbs/30_dims.ipynb 49
def vidpxs(mov: vidcap | pathlike) -> rect:
    '''
    Get the width and height (in pixels) of a video.

    Parameters
    ----------
    mov : VidCap | pathlike
        Video capture object or path to the video file.

    Returns
    -------
    Rect
        The width and height of the video in pixels.

    Examples
    --------
    >>> vidpxs('/path/to/video.avi')  # Returns the width and height of the video
    (1920, 1080)

    >>> vidpxs(cv2.VideoCapture('/path/to/video.avi')) 
    (1920, 1080)
    '''
    try: 
        import cv2
        CPFW, CPFH = cv2.CAP_PROP_FRAME_WIDTH, cv2.CAP_PROP_FRAME_HEIGHT
    except: 
        CPFW, CPFH = 3, 4
    if ispathlike(mov): mov = cv2.VideoCapture(str(mov))
    pxs = int(mov.get(CPFW)), (int(mov.get(CPFH)))
    mov.release()
    return pxs

# %% ../nbs/30_dims.ipynb 51
def subpxs(width: int, height: int, rows: int, cols: int) -> rect:
    '''
    Calculate the pixel dimensions for each subframe in a grid layout.

    Parameters
    ----------
    width : int
        Total width of the video or image.
    height : int
        Total height of the video or image.
    rows : int
        Number of rows in the grid.
    cols : int
        Number of columns in the grid.

    Returns
    -------
    Rect
        The width and height of each subframe in the grid.

    Examples
    --------
    >>> subpxs(1920, 1080, 2, 3)  # Subframe dimensions in a 2x3 grid
    (640, 540)
    '''
    return (width // cols, height // rows)

# %% ../nbs/30_dims.ipynb 54
def rbatches_dim(t: tensor, batch_first: bool = True) -> intq: 
    '''Extract the batch size of tensor `t` for recurrent neural networks.
    
    This will be either `None`, `x.size(0)` or `x.size(1)` depending on the 
    shape of `t` and the value of `batch_first`.

    Parameters
    ----------
    t : tensor
        The tensor from which to extract the batch size.
    
    batch_first : bool, default: True
        A flag indicating whether the batch dimension is the first dimension.
        
    Returns
    -------
    int or None
        The batch size of the tensor `t` if it is batched, otherwise `None`.

    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)

# %% ../nbs/30_dims.ipynb 56
def rchannels_dim(t: tensor, batch_first: bool = True) -> intq:
    '''Extract the number of channels in tensor `t` for recurrent neural networks.
    
    This will be either `x.size(0)` or `x.size(1)` depending on the 
    shape of `t` and the value of `batch_first`.
    
    Parameters
    ----------
    t : tensor
        The tensor from which to extract the batch size.
    
    batch_first : bool, default: True
        A flag indicating whether the batch dimension is the first dimension.
        
    Returns
    -------
    int or None
        The batch size of the tensor `t` if it is batched, otherwise `None`.
    
    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/30_dims.ipynb 58
def ndirections(bidirectional: bool = False) -> int:    
    '''
    Calculate the number of directions in a recurrent neural network layer.

    Parameters
    ----------
    bidirectional : bool, optional
        Indicates whether the layer is bidirectional, defaults to False.

    Returns
    -------
    int
        The number of directions (1 for unidirectional, 2 for bidirectional).

    Examples
    --------
    >>> directions(True)
    2
    >>> directions(False)
    1
    '''
    return int(2 if bidirectional else 1)

# %% ../nbs/30_dims.ipynb 60
def ndirected(nlays: int = 1, bidirectional: bool = False) -> int:
    '''
    Calculate the total number of layer-direction combinations in a recurrent neural network
    i.e. (D * L) where D is the number of directions and L is the number of layers.

    Parameters
    ----------
    nlays : int, optional
        The number of layers in the RNN, defaults to 1.
        
    bidirectional : bool, optional
        Indicates whether the layers are bidirectional, defaults to False.

    Returns
    -------
    int
        The total number of layer-direction combinations.

    Examples
    --------
    >>> directed_layers(2, True)
    4
    '''
    return nlays * ndirections(bidirectional)

# %% ../nbs/30_dims.ipynb 62
#| export
