# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/70_chns.ipynb.

# %% auto 0
__all__ = ['match_channels_format', 'ascfmt', 'as_channels_format', 'to_channels_format', 'aschnfmt', 'format_channels', 'aschns',
           'channels_resize', 'aschannels']

# %% ../nbs/70_chns.ipynb 6
from functools import wraps

# %% ../nbs/70_chns.ipynb 8
#| export


# %% ../nbs/70_chns.ipynb 10
try: import numpy as np
except: ...

# %% ../nbs/70_chns.ipynb 12
from quac import (tensor, int4, intq, size as Size)
from chck import isnone, isstr, istuple

# %% ../nbs/70_chns.ipynb 14
from asto.cons import (
    THWC, TCHW, THWC_AXES, TCHW_AXES,
    WRAPS_ASSIGN_ANNDOCS
)
from .atyp import P, Channels
from .tens import resize_tensor

# %% ../nbs/70_chns.ipynb 18
def _flip_channels_format(t: tensor) -> tensor:
    '''Flips the tensor between THWC and TCHW channel configurations.
    
    See Also
    --------
    etrc.enum.Channels.flip: Alias for _flip_channels_format
    '''
    return t.transpose(1, 3)

# %% ../nbs/70_chns.ipynb 20
def _guess_channels_index(t: tensor) -> int:
    '''Returns the index of the channel dimension
    
    See Also
    --------
    etrc.enum.Channels.cidx: Alias for _guess_channels_index
    '''
    idx = -1
    for i in range(len(t.shape)):
        if t.shape[i] != 3: 
            continue
        idx = i
    idx = idx if idx >= 0 else np.argmin(t.shape)
    if (val := t.shape[idx]) not in {1, 3}: 
        raise ValueError(f'Expected eitehr 1 or 3 channels, got {val}')
    return idx

# %% ../nbs/70_chns.ipynb 23
def _match_channels_format_string(fmt: str) -> str:
    match fmt.casefold():
        case THWC.casefold():
            return THWC
        case TCHW.casefold(): 
            return TCHW
        case _: 
            return THWC

# %% ../nbs/70_chns.ipynb 25
def _match_channels_format_tuple(fmt: int4) -> str:
    match fmt:
        case v if fmt == THWC_AXES: 
            return THWC
        case v if fmt == TCHW_AXES: 
            return TCHW
        case _: 
            return TCHW

# %% ../nbs/70_chns.ipynb 27
def match_channels_format(fmt: str | int4) -> str:
    '''Safely retrieves the channel format string, defaulting to 'THWC' if invalid.

    Parameters
    ----------
    fmt : Channels | 'thcw' | 'tchw' | (0, 1, 2, 3) | (0, 3, 1, 2)
        The channel configuration to validate.

    Returns
    -------
    Channels
        The validated channel configuration, defaults to THWC if input is invalid.
        
    See Also
    --------
    etrc.enum.Channels.safe: Alias for match_channels_format
    '''
    match fmt:
        case t if isstr(fmt): 
            return _match_channels_format_string(fmt)
        case t if istuple(fmt): 
            return _match_channels_format_tuple(fmt)
        case _: 
            return THWC

# %% ../nbs/70_chns.ipynb 29
def ascfmt(
    t: tensor, 
    fmt: Channels = THWC, 
    cidx: intq = None,
    *args: P.args,
    **kwargs: P.kwargs, 
) -> tensor:
    '''
    Converts a tensor to the specified channel configuration.

    Parameters
    ----------
    t : tensor
        The input tensor to convert.
    
    fmt : Channels | 'thcw' | 'tchw' | (0, 1, 2, 3) | (0, 3, 1, 2), default: 'thwc'
        The desired channel configuration to return.
        
    cidx : int, optional
        The current channel index in the tensor, calculated if not provided.
        
    Returns
    -------
    tensor
        The tensor converted to the specified channel configuration.
        
    See Also
    --------
    aschnfmt : Alias for ascfmt
    as_channels_format : Alias for ascfmt
    format_channels : Alias for ascfmt
    to_channels_format : Alias for ascfmt 
    etrc.enum.Channels.to : Alias for ascfmt
    '''
    # match the provided format to a valid channel configuration
    fmt = match_channels_format(fmt)
    # if channel_index is not provided, guess which axis belongs to channels 
    if isnone(cidx): cidx = _guess_channels_index(t)
    # flip the channels axis if needed
    match (fmt, cidx):
        case (THWC, 1): return _flip_channels_format(t)
        case (THWC, 3): return t
        case (TCHW, 1): return t
        case (TCHW, 3): return _flip_channels_format(t)
        case _: return t
    # raise ValueError(f'Invalid channel configuration: {fmt}, {channel_index}')

@wraps(ascfmt, assigned=WRAPS_ASSIGN_ANNDOCS)
def as_channels_format(t: tensor, fmt: Channels = THWC, cidx: intq = None, *args: P.args, **kwargs: P.kwargs) -> tensor:
    return ascfmt(t, fmt, cidx)

@wraps(ascfmt, assigned=WRAPS_ASSIGN_ANNDOCS)
def to_channels_format(t: tensor, fmt: Channels = THWC, cidx: intq = None, *args: P.args, **kwargs: P.kwargs) -> tensor:
    return ascfmt(t, fmt, cidx)

@wraps(ascfmt, assigned=WRAPS_ASSIGN_ANNDOCS)
def aschnfmt(t: tensor, fmt: Channels = THWC, cidx: intq = None, *args: P.args, **kwargs: P.kwargs) -> tensor:
    return ascfmt(t, fmt, cidx)

@wraps(ascfmt, assigned=WRAPS_ASSIGN_ANNDOCS)
def format_channels(t: tensor, fmt: Channels = THWC, cidx: intq = None, *args: P.args, **kwargs: P.kwargs) -> tensor:
    return ascfmt(t, fmt, cidx)

# %% ../nbs/70_chns.ipynb 31
def aschns(
    t: tensor, 
    size: Size | None = None, 
    fmt: Channels = THWC, 
    *args: P.args, 
    **kwargs: P.kwargs, 
) -> tensor:
    '''Converts a tensor to the specified channel configuration adjusting the size
    as needed.

    Parameters
    ----------
    t : tensor
        The input tensor to convert.
            
    size : tuple[int, ...], optional
        The desired size of the output tensor.
        
    fmt : Channels | 'thcw' | 'tchw' | (0, 1, 2, 3) | (0, 3, 1, 2), default: 'thwc'
        The desired channel format to return.
    
    Returns
    -------
    tensor
        The tensor converted to the specified channel configuration.
        
    See Also
    --------
    aschannels : Alias for aschns
    channels_resize : Alias for aschns
    etrc.enum.Channels.resize : Alias for aschns
    '''
    if isnone(size): return as_channels_format(t, fmt=fmt)
    # Cast to TCHW format for compatiability with `torchvision.transforms.functional.resize`
    t = as_channels_format(t, fmt=TCHW)
    # Then we apply the `resize` operation
    t = resize_tensor(t, size)
    # Then we convert back to the desired format
    t = as_channels_format(t, fmt=fmt)
    return t

@wraps(aschns, assigned=WRAPS_ASSIGN_ANNDOCS)
def channels_resize(t: tensor, size: Size | None = None, fmt: Channels = THWC, *args: P.args, **kwargs: P.kwargs) -> tensor:
    return aschns(t, size=size, fmt=fmt, *args, **kwargs)

@wraps(aschns, assigned=WRAPS_ASSIGN_ANNDOCS)
def aschannels(t: tensor, size: Size | None = None, fmt: Channels = THWC, *args: P.args, **kwargs: P.kwargs) -> tensor:
    return aschns(t, size=size, fmt=fmt, *args, **kwargs)
