# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/00_vision_plot.ipynb.

# %% auto 0
__all__ = ['validate_anatomical_plane', 'show_med_img', 'find_max_slice']

# %% ../nbs/00_vision_plot.ipynb 1
from fastai.data.all import *
from torchio.visualization import rotate

# %% ../nbs/00_vision_plot.ipynb 3
def _get_slice(image, channel: int, slice_index: int, anatomical_plane: int, voxel_size: (int, list)):
    """
    A private method to get a 2D tensor and aspect ratio for plotting.
    This is modified code from the torchio function `plot_volume`.

    Args:
        image: The input image.
        channel: Channel of the image.
        slice_index: Index of the 2D slice.
        anatomical_plane: Anatomical plane of the image.
        voxel_size: Voxel size for the image.

    Returns:
        A 2D tensor containing the sliced image and its aspect.
    """
    if voxel_size is None:
        warnings.warn('Voxel size not defined. Aspect ratio of the plot might not be correct.')
        voxel_size = 1

    if not is_listy(voxel_size):
        voxel_size = (voxel_size, voxel_size, voxel_size)

    sr, sa, ss = voxel_size
    data = image.data[channel]

    if slice_index is None:
        slice_center = np.array(data.shape) // 2
        slice_index = slice_center[anatomical_plane]

    operations = {
        0: (lambda d, i: d[i, :, :], ss / sa),
        1: (lambda d, i: d[:, i, :], ss / sr),
        2: (lambda d, i: d[:, :, i], sa / sr)
    }

    sliced_function, aspect_ratio = operations.get(anatomical_plane)
    sliced_img = rotate(sliced_function(data, slice_index), radiological=True, n=1)

    return sliced_img, aspect_ratio

# %% ../nbs/00_vision_plot.ipynb 4
def validate_anatomical_plane(anatomical_plane):
    """Ensure anatomical_plane is either 0, 1, or 2."""
    if anatomical_plane not in [0, 1, 2]:
        raise ValueError("anatomical_plane should be either 0, 1, or 2.")

# %% ../nbs/00_vision_plot.ipynb 5
@delegates(plt.Axes.imshow, keep=True, but=['shape', 'imlim'])
def show_med_img(im, ctx, channel: int, slice_index: int, anatomical_plane: int,
                 voxel_size: (int, list), ax=None, figsize=None, title=None, **kwargs):
    """Show an image on `ax`. This is a modified code from the fastai function `show_image`.

    Args:
        im: The input image.
        ctx: The context.
        channel: Channel of the image.
        slice_index: Index of the 2D slice.
        anatomical_plane: Anatomical plane of the image.
        voxel_size: Voxel size for the image.
        ax: Axis for the plot.
        figsize: Figure size for the plot.
        title: Title for the plot.
        kwargs: Additional parameters for plt.Axes.imshow method.

    Returns:
        Axis with the plot.
    """
                     
    validate_anatomical_plane(anatomical_plane)
                     
    if hasattrs(im, ('data', 'cpu', 'permute')): # Check if `im` has the necessary attributes
        im = im.data.cpu()
        im, aspect = _get_slice(
            im, 
            channel=channel, 
            anatomical_plane=anatomical_plane,
            voxel_size=voxel_size, 
            slice_index=slice_index
        )

    ax = ax if ax is not None else ctx 

    if ax is None: # ax is only None when .show() is used.
        _, ax = plt.subplots(figsize=figsize)

    ax.imshow(im, aspect=aspect, **kwargs)

    if title is not None:
        ax.set_title(title)

    ax.axis('off')

    return ax

# %% ../nbs/00_vision_plot.ipynb 6
def find_max_slice(mask_data, anatomical_plane):
    """Find slice index based on mask"""
    
    axis_mapping = {0: (1, 2),
                    1: (0, 2),
                    2: (0, 1)}

    sums = mask_data.sum(axis=axis_mapping.get(anatomical_plane, (1, 2)))
    idx = np.argmax(sums)

    return idx
