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

# %% auto 0
__all__ = ['match_channels_format', 'to_channels_format', 'ascfmt', 'aschnfmt']

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

# %% ../nbs/70_chns.ipynb 8
from types import ModuleType
from typing import Tuple

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

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

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

# %% ../nbs/70_chns.ipynb 17
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 19
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 21
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 23
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 25
def match_channels_format(fmt: str | int4) -> str:
    '''Safely retrieves the channel format string, defaulting to 'THWC' if invalid.

    Parameters
    ----------
    fmt : '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 27
def to_channels_format(t: tensor, fmt: str | int4 = THWC, channel_index: intq = None) -> tensor:
    '''
    Converts a tensor to the specified channel configuration.

    Parameters
    ----------
    t : Tensor
        The input tensor to convert.
    c : Channels, optional
        The target channel configuration, defaults to THWC.
    cidx : int, optional
        The current channel index in the tensor, calculated if not provided.

    Returns
    -------
    Tensor
        The tensor converted to the specified channel configuration.
    '''
    fmt = match_channels_format(fmt)
    if isnone(channel_index): 
        channel_index = _guess_channels_index(t)
        
    match (fmt, channel_index):
        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(to_channels_format)
def ascfmt(t: tensor, fmt: str | int4 = THWC, channel_index: intq = None) -> tensor:
    return to_channels_format(t, fmt, channel_index)

@wraps(to_channels_format)
def aschnfmt(t: tensor, fmt: str | int4 = THWC, channel_index: intq = None) -> tensor:
    return to_channels_format(t, fmt, channel_index)
