reshape_arrays

def get_original_stitch_slices(
    array_shape: Sequence[int],
    original_axes: str,
    original_shape: Sequence[int],
    sample_idx: int,
    stitch_coords: Sequence[int],
    crop_size: Sequence[int],
) -> tuple[slice | int, ...]:
    """Get slices to stitch tile back into original array.

    `sample_idx and `stitch_coords` are expressed with respect to the transofrmed space
    (SCZYX or SCYX). The returned slices will index into the original array for
    stitching the tile back in place.

    Parameters
    ----------
    array_shape : Sequence[int]
        Shape of the array in transformed space (e.g. SCZYX).
    original_axes : str
        Original axes string of the full data.
    original_shape : Sequence[int]
        Original shape of the full data.
    sample_idx : int
        Index of the sample in transformed space (S axis) to stitch back.
    stitch_coords : Sequence[int]
        Starting coordinates of the tile in the original spatial axes (Y, X and Z if
        present).
    crop_size : Sequence[int]
        Size of the tile in the original spatial axes (Y, X and Z if present).

    Returns
    -------
    tuple[slice | int, ...]
        Slices to index into the original array for stitching the tile back in place.
    """
    transform = AxesTransform(original_axes, tuple(original_shape))

    stitch_slices: list[slice | int] = []
    which_axes = []

    # unravel sample indices
    if len(transform.sample_dims) >= 1:
        sample_dims = transform.sample_dims
        sample_dim_sizes = [transform.dim_sizes[d] for d in sample_dims]
        sample_indices = [
            int(i) for i in np.unravel_index(sample_idx, sample_dim_sizes)
        ]
        stitch_slices.extend(sample_indices)
        which_axes.extend(sample_dims)

    if not transform.c_added:
        stitch_slices.append(slice(0, array_shape[1]))
        which_axes.append("C")

    # add spatial slices
    stitch_slices.extend(
        [
            slice(start, start + length)
            for start, length in zip(stitch_coords, crop_size, strict=True)
        ]
    )
    which_axes.extend([a for a in transform.transformed_axes if a in "ZYX"])
    assert len(stitch_slices) == len(transform.original_axes)

    # reorder slices
    stitch_slices = [
        stitch_slices[which_axes.index(a)] for a in transform.original_axes
    ]

    return tuple(stitch_slices)

def reshape_array(
    array: NDArray,
    original_axes: str,
    original_shape: Sequence[int],
) -> NDArray:
    """Reshape array from arbitrary axes order to `SC(Z)YX`.

    Parameters
    ----------
    array : numpy.ndarray
        Input array.
    original_axes : str
        Original axes string describing current dimension order (e.g. `YXC`).
    original_shape : Sequence[int]
        Original shape of the array.

    Returns
    -------
    numpy.ndarray
        Array reshaped to `SC(Z)YX`.
    """
    transform = AxesTransform(original_axes, original_shape)

    # reorder axes to reference STCZYX
    permutation = transform.order_permutation
    array = np.moveaxis(array, permutation, list(range(len(permutation))))

    # merge sample dims
    n_sample = len(transform.sample_dims)
    if n_sample > 1:
        array = array.reshape((-1,) + array.shape[n_sample:])
    elif n_sample == 0:
        array = array[np.newaxis]

    # add singleton C, if necessary
    if transform.c_added:
        array = np.expand_dims(array, 1)

    return array

def restore_array(
    array: NDArray,
    original_axes: str,
    original_shape: Sequence[int],
) -> NDArray:
    """Restore array from `SC(Z)YX` space back to original axes and shape.

    Parameters
    ----------
    array : numpy.ndarray
        Array in `SC(Z)YX` format.
    original_axes : str
        Original axes string (e.g. `YXC`).
    original_shape : Sequence[int]
        Original shape of the data.

    Returns
    -------
    numpy.ndarray
        Array with original axes order and shape restored.

    Raises
    ------
    ValueError
        If input array is not 4D (SCYX) or 5D (SCZYX), or if restoring shape is not
        supported for the given original axes (e.g. T as Z with CZI format).
    """
    if len(array.shape) not in (4, 5):
        raise ValueError(f"Expected 4D (SCYX) or 5D (SCZYX), got {len(array.shape)}D.")

    if len(array.shape) == 5 and original_axes == "SCTYX":
        raise ValueError(
            "Restoring shape is currently not supported for CZI format (T used as "
            "depth axis)."
        )

    transform = AxesTransform(original_axes, tuple(original_shape))
    current_axes = list(transform.transformed_axes)
    sample_dims = transform.sample_dims

    # restore sample dimensions
    # - if multiple sample dims, it will reshape the array and add all sample dims back
    # - if single sample dim that is not S, it will rename the dimensions
    # - if no sample dims, it will remove the singleton S dim
    sizes = tuple(transform.dim_sizes[d] for d in sample_dims)
    array = array.reshape(sizes + array.shape[1:])
    current_axes = list(sample_dims) + current_axes[1:]

    # remove singleton C
    if transform.c_added:
        c_idx = current_axes.index("C")
        array = np.squeeze(array, axis=c_idx)
        current_axes.pop(c_idx)

    # reorder axes to original order
    current_str = "".join(current_axes)
    if current_str != original_axes:
        source = [current_str.index(a) for a in original_axes]
        array = np.moveaxis(array, source, list(range(len(original_axes))))

    return array

def restore_tile(
    tile: NDArray,
    original_axes: str,
    original_shape: Sequence[int],
) -> NDArray:
    """Restore single tile from `C(Z)YX` space back to original axes and shape.

    Parameters
    ----------
    tile : numpy.ndarray
        Tile in `C(Z)YX` format.
    original_axes : str
        Original axes string of the full data.
    original_shape : Sequence[int]
        Original shape of the full data.

    Returns
    -------
    numpy.ndarray
        Tile with original spatial axes order restored.
    """
    if len(tile.shape) not in (3, 4):
        raise ValueError(f"Expected 3D (CYX) or 4D (CZYX), got {len(tile.shape)}D.")

    transform = AxesTransform(original_axes, tuple(original_shape))

    # remove sample dim
    current_axes = list(transform.transformed_axes[1:])

    # remove singleton C if it was added
    if transform.c_added:
        tile = np.squeeze(tile, axis=0)
        current_axes.pop(0)

    # expected dimensions for a tile in original space
    tile_dims = set(transform.transformed_axes[1:])
    tile_original = "".join(a for a in original_axes if a in tile_dims)

    current_str = "".join(current_axes)
    if current_str != tile_original:
        source = [current_str.index(a) for a in tile_original]
        tile = np.moveaxis(tile, source, list(range(len(tile_original))))

    return tile