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

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

# %% ../nbs/00_vision_plot.ipynb #7ea97ee7-8a55-485d-a75c-ae1a8055f489
import warnings
from fastai.data.all import *
from torchio.visualization import rotate

# %% ../nbs/00_vision_plot.ipynb #55a13906-aef5-445a-8919-eb7daa3816ea
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 #4a0a8991-fbb8-4d07-80b6-804852139253
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 #7955e15b-7580-4219-838e-93ff094e146a
@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 #c31d23c6-d819-4039-8d53-ea36ef20f049
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

# %% ../nbs/00_vision_plot.ipynb #ah52hmsl0of
def show_segmentation_comparison(
    image, ground_truth, prediction,
    slice_index: int = None,
    anatomical_plane: int = 2,
    metric_value: float = None,
    metric_name: str = 'DSC',
    channel: int = 0,
    figsize: tuple = (15, 5),
    cmap_img: str = 'gray',
    cmap_mask: str = 'gray',
    voxel_size = None
):
    """Display 3-panel comparison: Image | Ground Truth | Prediction.
    
    Useful for validating segmentation results, especially after patch-based 
    inference where results are not in standard fastai batch format.
    
    Args:
        image: Input image (MedImage, MedMask, or tensor [C, H, W, D])
        ground_truth: Ground truth mask (MedMask or tensor [C, H, W, D])
        prediction: Predicted mask (tensor [C, H, W, D])
        slice_index: Slice to display. If None, uses find_max_slice on ground_truth.
        anatomical_plane: 0=sagittal, 1=coronal, 2=axial (default)
        metric_value: Optional metric value to display in prediction title
        metric_name: Name of metric for title (default 'DSC')
        channel: Channel to display for multi-channel data (default 0)
        figsize: Figure size (default (15, 5))
        cmap_img: Colormap for image (default 'gray')
        cmap_mask: Colormap for masks (default 'gray')
        voxel_size: Voxel spacing for aspect ratio. If None, aspect=1.

    Example::
    
        # After patch_inference()
        show_segmentation_comparison(
            image=val_img,
            ground_truth=val_gt,
            prediction=predictions[0],
            metric_value=results_df.iloc[0]['DSC'],
            anatomical_plane=2
        )
    """
    validate_anatomical_plane(anatomical_plane)
    
    # Ensure we have tensor data
    img_data = image.data if hasattr(image, 'data') else image
    gt_data = ground_truth.data if hasattr(ground_truth, 'data') else ground_truth
    pred_data = prediction.data if hasattr(prediction, 'data') else prediction
    
    # Move to CPU if needed
    if hasattr(img_data, 'cpu'): img_data = img_data.cpu()
    if hasattr(gt_data, 'cpu'): gt_data = gt_data.cpu()
    if hasattr(pred_data, 'cpu'): pred_data = pred_data.cpu()
    
    # Find optimal slice if not provided
    if slice_index is None:
        gt_np = gt_data[channel].numpy()
        slice_index = find_max_slice(gt_np, anatomical_plane)
    
    # Create figure
    fig, axes = plt.subplots(1, 3, figsize=figsize)
    
    # Get slices with proper aspect ratio using existing _get_slice helper
    img_slice, img_aspect = _get_slice(img_data, channel, slice_index, anatomical_plane, voxel_size)
    gt_slice, gt_aspect = _get_slice(gt_data, channel, slice_index, anatomical_plane, voxel_size)
    pred_slice, pred_aspect = _get_slice(pred_data, channel, slice_index, anatomical_plane, voxel_size)
    
    # Plot panels
    axes[0].imshow(img_slice, cmap=cmap_img, aspect=img_aspect)
    axes[0].set_title('Input Image')
    axes[0].axis('off')
    
    axes[1].imshow(gt_slice, cmap=cmap_mask, aspect=gt_aspect)
    axes[1].set_title('Ground Truth')
    axes[1].axis('off')
    
    # Build prediction title
    pred_title = 'Prediction'
    if metric_value is not None:
        pred_title = f'Prediction ({metric_name}: {metric_value:.4f})'
    
    axes[2].imshow(pred_slice, cmap=cmap_mask, aspect=pred_aspect)
    axes[2].set_title(pred_title)
    axes[2].axis('off')
    
    plt.tight_layout()
