# AUTOGENERATED! DO NOT EDIT! File to edit: ../../nbs/utils/06_rearrange.ipynb.

# %% auto 0
__all__ = ['make_channels_first', 'make_channels_last', 'rearrange_channels']

# %% ../../nbs/utils/06_rearrange.ipynb 2
import numpy as np
import torch
from einops import rearrange

# %% ../../nbs/utils/06_rearrange.ipynb 4
def make_channels_first(x: torch.Tensor | np.ndarray):
    """Convert an n-dimensional tensor or array to channels first format.

    Args:
        x: The input tensor / array. Should have at least 3 dimensions.

    Returns:
        The input tensor / array in channels first format.
    """
    x = rearrange(x, "b ... d -> b d ...")
    if torch.is_tensor(x):
        x = x.contiguous()

    return x

# %% ../../nbs/utils/06_rearrange.ipynb 6
def make_channels_last(x: torch.Tensor | np.ndarray):
    """Convert an n-dimensional tensor or array to channels last format.

    Args:
        x: The input tensor / array. Should have at least 3 dimensions.

    Returns:
        The input tensor / array in channels last format.
    """
    x = rearrange(x, "b d ... -> b ... d")
    if torch.is_tensor(x):
        x = x.contiguous()

    return x

# %% ../../nbs/utils/06_rearrange.ipynb 8
def rearrange_channels(x: torch.Tensor | np.ndarray, cur_channels_first: bool, new_channels_first: bool):
    """Rearrange the channels of a tensor / array to either channels_first or channels_last format.

    Args:
        x: The input tensor / array.
        cur_channels_first: Whether the input tensor / array is in channels first format.
        new_channels_first: Whether the output should be in channels first format.

    Returns:
        The input tensor / array with the channels rearranged.
    """

    if cur_channels_first is new_channels_first:
        return x
    elif cur_channels_first:
        return make_channels_last(x)
    else:
        return make_channels_first(x)
