# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/03_vision_augment.ipynb.

# %% auto #0
__all__ = ['CustomDictTransform', 'do_pad_or_crop', 'PadOrCrop', 'ZNormalization', 'RescaleIntensity', 'NormalizeIntensity',
           'BraTSMaskConverter', 'BinaryConverter', 'RandomGhosting', 'RandomSpike', 'RandomNoise', 'RandomBiasField',
           'RandomBlur', 'RandomGamma', 'RandomIntensityScale', 'RandomMotion', 'RandomAnisotropy', 'RandomCutout',
           'RandomElasticDeformation', 'RandomAffine', 'RandomFlip', 'OneOf', 'GpuPatchAugmentation',
           'gpu_patch_augmentations', 'suggest_patch_augmentations']

# %% ../nbs/03_vision_augment.ipynb #2d6694aa
from fastai.data.all import *
from .vision_core import *
import torch.nn.functional as F
import math
import torchio as tio
from monai.transforms import NormalizeIntensity as MonaiNormalizeIntensity

# %% ../nbs/03_vision_augment.ipynb #7c2bacd9
class CustomDictTransform(ItemTransform):
    """A class that serves as a wrapper to perform an identical transformation on both
    the image and the target (if it's a mask).
    """

    split_idx = 0  # Only perform transformations on training data. Use TTA() for transformations on validation data.

    def __init__(self, aug):
        """Constructs CustomDictTransform object.

        Args:
            aug (Callable): Function to apply augmentation on the image.
        """
        self.aug = aug

    @property
    def tio_transform(self):
        """Return the underlying TorchIO transform.

        This property enables using fastMONAI wrappers in patch-based workflows
        where raw TorchIO transforms are needed for tio.Compose().
        """
        return self.aug

    def encodes(self, x):
        """
        Applies the stored transformation to an image, and the same random transformation
        to the target if it is a mask. If the target is not a mask, it returns the target as is.

        Args:
            x (Tuple[MedImage, Union[MedMask, TensorCategory]]): A tuple containing the
            image and the target.

        Returns:
            Tuple[MedImage, Union[MedMask, TensorCategory]]: The transformed image and target.
            If the target is a mask, it's transformed identically to the image. If the target
            is not a mask, the original target is returned.
        """
        img, y_true = x
        
        # Use identity affine if MedImage.affine_matrix is not set
        affine = MedImage.affine_matrix if MedImage.affine_matrix is not None else np.eye(4)

        if isinstance(y_true, (MedMask)):
            aug = self.aug(tio.Subject(img=tio.ScalarImage(tensor=img, affine=affine),
                                        mask=tio.LabelMap(tensor=y_true, affine=affine)))
            return MedImage.create(aug['img'].data), MedMask.create(aug['mask'].data)

        aug = self.aug(tio.Subject(img=tio.ScalarImage(tensor=img)))
        return MedImage.create(aug['img'].data), y_true

# %% ../nbs/03_vision_augment.ipynb #d7e1fdf2
def do_pad_or_crop(o, target_shape, padding_mode, mask_name, dtype=torch.Tensor):
    #TODO:refactorize
    pad_or_crop = tio.CropOrPad(target_shape=target_shape, padding_mode=padding_mode, mask_name=mask_name)
    return dtype(pad_or_crop(o))

# %% ../nbs/03_vision_augment.ipynb #9c75bb84
class PadOrCrop(DisplayedTransform):
    """Resize image using TorchIO `CropOrPad`."""
    
    order = 0

    def __init__(self, size, padding_mode=0, mask_name=None):
        if not is_listy(size): 
            size = [size, size, size]
        self.pad_or_crop = tio.CropOrPad(target_shape=size,
                                    padding_mode=padding_mode, 
                                    mask_name=mask_name)

    @property
    def tio_transform(self):
        """Return the underlying TorchIO transform."""
        return self.pad_or_crop

    def encodes(self, o: (MedImage, MedMask)):
        return type(o)(self.pad_or_crop(o))

# %% ../nbs/03_vision_augment.ipynb #534509q2nn
def _foreground_masking(tensor):
    """Mask for non-zero voxels (nnU-Net-style foreground normalization)."""
    return tensor > 0

# %% ../nbs/03_vision_augment.ipynb #ca95a690
class ZNormalization(DisplayedTransform):
    """Apply TorchIO `ZNormalization`."""

    order = 0

    def __init__(self, masking_method=None, channel_wise=True):
        if masking_method == 'foreground':
            masking_method = _foreground_masking
        self.z_normalization = tio.ZNormalization(masking_method=masking_method)
        self.channel_wise = channel_wise

    @property
    def tio_transform(self):
        """Return the underlying TorchIO transform."""
        return self.z_normalization

    def encodes(self, o: MedImage):
        try:
            if self.channel_wise:
                o = torch.stack([self.z_normalization(c[None])[0] for c in o])
            else: 
                o = self.z_normalization(o)
        except RuntimeError as e:
            if "Standard deviation is 0" in str(e):
                # Calculate mean for debugging information
                mean = float(o.mean())
                
                error_msg = (
                    f"Standard deviation is 0 for image (mean={mean:.3f}).\n"
                    f"This indicates uniform pixel values.\n\n"
                    f"Possible causes:\n"
                    f"• Corrupted or blank image\n"
                    f"• Oversaturated regions\n"  
                    f"• Background-only regions\n"
                    f"• All-zero mask being processed as image\n\n"
                    f"Suggested solutions:\n"
                    f"• Check image quality and acquisition\n"
                    f"• Verify image vs mask data loading"
                )
                raise RuntimeError(error_msg) from e

        return MedImage.create(o)

    def encodes(self, o: MedMask):
        return o

# %% ../nbs/03_vision_augment.ipynb #ee9cf2a3
class RescaleIntensity(DisplayedTransform):
    """Apply TorchIO RescaleIntensity for robust intensity scaling.
    
    Args:
        out_min_max (tuple[float, float]): Output intensity range (min, max)
        in_min_max (tuple[float, float]): Input intensity range (min, max) 
    
    Example for CT images:
        # Normalize CT from air (-1000 HU) to bone (1000 HU) into range (-1, 1)
        transform = RescaleIntensity(out_min_max=(-1, 1), in_min_max=(-1000, 1000))
    """
    
    order = 0
    
    def __init__(self, out_min_max: tuple[float, float], in_min_max: tuple[float, float]):
        self.rescale = tio.RescaleIntensity(out_min_max=out_min_max, in_min_max=in_min_max)

    @property
    def tio_transform(self):
        """Return the underlying TorchIO transform."""
        return self.rescale

    def encodes(self, o: MedImage):
        return MedImage.create(self.rescale(o))
        
    def encodes(self, o: MedMask):
        return o

# %% ../nbs/03_vision_augment.ipynb #341dc45f
class _TioNormalizeIntensity(tio.IntensityTransform):
    """TorchIO-compatible wrapper for MONAI NormalizeIntensity.

    Enables NormalizeIntensity to work in patch-based workflows
    where raw TorchIO transforms are needed for tio.Compose().
    """

    def __init__(self, nonzero=True, channel_wise=True,
                 subtrahend=None, divisor=None, **kwargs):
        super().__init__(p=1, **kwargs)
        self.channel_wise = channel_wise
        self.transform = MonaiNormalizeIntensity(
            nonzero=nonzero,
            channel_wise=False,
            subtrahend=subtrahend,
            divisor=divisor
        )

    def apply_transform(self, subject):
        for image in self.get_images(subject):
            data = image.data
            if self.channel_wise:
                result = torch.stack([self.transform(c[None])[0] for c in data])
            else:
                result = torch.Tensor(self.transform(data))
            image.set_data(result)
        return subject


class NormalizeIntensity(DisplayedTransform):
    """Apply MONAI NormalizeIntensity.
    
    Args:
        nonzero (bool): Only normalize non-zero values (default: True)
        channel_wise (bool): Apply normalization per channel (default: True)
        subtrahend (float, optional): Value to subtract  
        divisor (float, optional): Value to divide by
    """
    
    order = 0
    
    def __init__(self, nonzero: bool = True, channel_wise: bool = True, 
                 subtrahend: float = None, divisor: float = None):
        self.nonzero = nonzero
        self.channel_wise = channel_wise
        self.subtrahend = subtrahend
        self.divisor = divisor
        
        self.transform = MonaiNormalizeIntensity(
            nonzero=nonzero,
            channel_wise=False,  # Always 'False', we handle channel-wise manually
            subtrahend=subtrahend,
            divisor=divisor
        )
        self._tio_normalize = _TioNormalizeIntensity(
            nonzero=nonzero,
            channel_wise=channel_wise,
            subtrahend=subtrahend,
            divisor=divisor
        )

    @property
    def tio_transform(self):
        """Return TorchIO-compatible transform for patch-based workflows."""
        return self._tio_normalize
    
    def encodes(self, o: MedImage):
        if self.channel_wise:
            result = torch.stack([self.transform(c[None])[0] for c in o])
        else:
            result = torch.Tensor(self.transform(o))
        
        return MedImage.create(result)
    
    def encodes(self, o: MedMask):
        return o

# %% ../nbs/03_vision_augment.ipynb #be431cc2
class BraTSMaskConverter(DisplayedTransform):
    '''Convert BraTS masks.'''

    order=1

    def encodes(self, o:(MedImage)): return o

    def encodes(self, o:(MedMask)):
        o = torch.where(o==4, 3., o)
        return MedMask.create(o)

# %% ../nbs/03_vision_augment.ipynb #f63701cb
class BinaryConverter(DisplayedTransform):
    '''Convert to binary mask.'''

    order=1

    def encodes(self, o: MedImage): 
        return o

    def encodes(self, o: MedMask):
        o = torch.where(o>0, 1., 0)
        return MedMask.create(o)

# %% ../nbs/03_vision_augment.ipynb #6b5a795f
class RandomGhosting(DisplayedTransform):
    """Apply TorchIO `RandomGhosting`."""
    
    split_idx, order = 0, 1

    def __init__(self, intensity=(0.5, 1), p=0.5):
        self.add_ghosts = tio.RandomGhosting(intensity=intensity, p=p)

    @property
    def tio_transform(self):
        """Return the underlying TorchIO transform."""
        return self.add_ghosts

    def encodes(self, o: MedImage):
        return MedImage.create(self.add_ghosts(o))

    def encodes(self, o: MedMask):
        return o

# %% ../nbs/03_vision_augment.ipynb #bf9ce6fc
class RandomSpike(DisplayedTransform):
    '''Apply TorchIO `RandomSpike`.'''
    
    split_idx, order = 0, 1

    def __init__(self, num_spikes=1, intensity=(1, 3), p=0.5):
        self.add_spikes = tio.RandomSpike(num_spikes=num_spikes, intensity=intensity, p=p)

    @property
    def tio_transform(self):
        """Return the underlying TorchIO transform."""
        return self.add_spikes

    def encodes(self, o: MedImage): 
        return MedImage.create(self.add_spikes(o))
        
    def encodes(self, o: MedMask):
        return o

# %% ../nbs/03_vision_augment.ipynb #a48e315a
class RandomNoise(DisplayedTransform):
    '''Apply TorchIO `RandomNoise`.'''

    split_idx, order = 0, 1

    def __init__(self, mean=0, std=(0, 0.25), p=0.5):
        self.add_noise = tio.RandomNoise(mean=mean, std=std, p=p)

    @property
    def tio_transform(self):
        """Return the underlying TorchIO transform."""
        return self.add_noise

    def encodes(self, o: MedImage): 
        return MedImage.create(self.add_noise(o))
    
    def encodes(self, o: MedMask):
        return o

# %% ../nbs/03_vision_augment.ipynb #6623a796
class RandomBiasField(DisplayedTransform):
    '''Apply TorchIO `RandomBiasField`.'''

    split_idx, order = 0, 1

    def __init__(self, coefficients=0.5, order=3, p=0.5):
        self.add_biasfield = tio.RandomBiasField(coefficients=coefficients, order=order, p=p)

    @property
    def tio_transform(self):
        """Return the underlying TorchIO transform."""
        return self.add_biasfield

    def encodes(self, o: MedImage): 
        return MedImage.create(self.add_biasfield(o))
        
    def encodes(self, o: MedMask):
        return o

# %% ../nbs/03_vision_augment.ipynb #7e7d29c1
class RandomBlur(DisplayedTransform):
    '''Apply TorchIO `RandomBlur`.'''

    split_idx, order = 0, 1

    def __init__(self, std=(0, 2), p=0.5):
        self.add_blur = tio.RandomBlur(std=std, p=p)

    @property
    def tio_transform(self):
        """Return the underlying TorchIO transform."""
        return self.add_blur

    def encodes(self, o: MedImage): 
        return MedImage.create(self.add_blur(o))
    
    def encodes(self, o: MedMask):
        return o

# %% ../nbs/03_vision_augment.ipynb #cd45c2c9
class RandomGamma(DisplayedTransform):
    '''Apply TorchIO `RandomGamma`.'''

    split_idx, order = 0, 1

    def __init__(self, log_gamma=(-0.3, 0.3), p=0.5):
        self.add_gamma = tio.RandomGamma(log_gamma=log_gamma, p=p)

    @property
    def tio_transform(self):
        """Return the underlying TorchIO transform."""
        return self.add_gamma

    def encodes(self, o: MedImage): 
        return MedImage.create(self.add_gamma(o))
    
    def encodes(self, o: MedMask):
        return o

# %% ../nbs/03_vision_augment.ipynb #4116aa33
class _TioRandomIntensityScale(tio.IntensityTransform):
    """TorchIO-compatible RandomIntensityScale for patch-based workflows.

    Randomly scales image intensities by a multiplicative factor.
    Only applies to ScalarImage keys (not LabelMap), which is the standard
    TorchIO IntensityTransform behavior.
    """

    def __init__(self, scale_range=(0.5, 2.0), p=0.5, **kwargs):
        super().__init__(p=p, **kwargs)
        self.scale_range = scale_range

    def apply_transform(self, subject):
        for image in self.get_images(subject):
            scale = torch.empty(1).uniform_(self.scale_range[0], self.scale_range[1]).item()
            image.set_data(image.data * scale)
        return subject


class RandomIntensityScale(DisplayedTransform):
    """Randomly scale image intensities by a multiplicative factor.

    Useful for domain generalization across different acquisition protocols
    with varying intensity ranges.

    Args:
        scale_range (tuple[float, float]): Range of scale factors (min, max).
            Values > 1 increase intensity, < 1 decrease intensity.
        p (float): Probability of applying the transform (default: 0.5)

    Example:
        # Scale intensities randomly between 0.5x and 2.0x
        transform = RandomIntensityScale(scale_range=(0.5, 2.0), p=0.3)
    """

    split_idx, order = 0, 1

    def __init__(self, scale_range: tuple[float, float] = (0.5, 2.0), p: float = 0.5):
        self.scale_range = scale_range
        self.p = p
        self._tio_intensity_scale = _TioRandomIntensityScale(
            scale_range=scale_range, p=p
        )

    @property
    def tio_transform(self):
        """Return TorchIO-compatible transform for patch-based workflows."""
        return self._tio_intensity_scale

    def encodes(self, o: MedImage):
        if torch.rand(1).item() > self.p:
            return o
        scale = torch.empty(1).uniform_(self.scale_range[0], self.scale_range[1]).item()
        return MedImage.create(o * scale)

    def encodes(self, o: MedMask):
        return o

# %% ../nbs/03_vision_augment.ipynb #bf89ffbe
class RandomMotion(DisplayedTransform):
    """Apply TorchIO `RandomMotion`."""

    split_idx, order = 0, 1

    def __init__(
            self, 
            degrees=10, 
            translation=10, 
            num_transforms=2, 
            image_interpolation='linear', 
            p=0.5
        ):
        self.add_motion = tio.RandomMotion(
            degrees=degrees, 
            translation=translation, 
            num_transforms=num_transforms, 
            image_interpolation=image_interpolation, 
            p=p
        )

    @property
    def tio_transform(self):
        """Return the underlying TorchIO transform."""
        return self.add_motion

    def encodes(self, o: MedImage):
        return MedImage.create(self.add_motion(o))

    def encodes(self, o: MedMask):
        return o

# %% ../nbs/03_vision_augment.ipynb #cl3ei8hm3z9
class RandomAnisotropy(DisplayedTransform):
    '''Apply TorchIO `RandomAnisotropy`.'''

    split_idx, order = 0, 1

    def __init__(self, axes=(0, 1, 2), downsampling=(1.5, 5),
                 image_interpolation='linear', scalars_only=True, p=0.5):
        self.add_anisotropy = tio.RandomAnisotropy(
            axes=axes, downsampling=downsampling,
            image_interpolation=image_interpolation,
            scalars_only=scalars_only, p=p)

    @property
    def tio_transform(self):
        """Return the underlying TorchIO transform."""
        return self.add_anisotropy

    def encodes(self, o: MedImage):
        return MedImage.create(self.add_anisotropy(o))

    def encodes(self, o: MedMask):
        return o

# %% ../nbs/03_vision_augment.ipynb #e7ea6486
def _create_ellipsoid_mask(shape, center, radii):
    """Create a 3D ellipsoid mask.

    Args:
        shape: (D, H, W) shape of the volume
        center: (z, y, x) center of ellipsoid
        radii: (rz, ry, rx) radii along each axis

    Returns:
        Boolean mask where True = inside ellipsoid
    """
    z, y, x = torch.meshgrid(
        torch.arange(shape[0]),
        torch.arange(shape[1]),
        torch.arange(shape[2]),
        indexing='ij'
    )
    dist = ((z - center[0]) / radii[0]) ** 2 + \
           ((y - center[1]) / radii[1]) ** 2 + \
           ((x - center[2]) / radii[2]) ** 2
    return dist <= 1.0

# %% ../nbs/03_vision_augment.ipynb #935c13e7
class _TioRandomCutout(tio.IntensityTransform):
    """TorchIO-compatible RandomCutout for patch-based workflows.

    When mask_only=True, cutouts only affect voxels where the mask is positive.
    The mask should be available in the Subject as 'mask' key.
    """

    def __init__(self, holes=1, spatial_size=8, fill_value=None,
                 max_holes=None, max_spatial_size=None, mask_only=True, p=0.2, **kwargs):
        super().__init__(p=p, **kwargs)
        self.holes = holes
        self.spatial_size = spatial_size
        self.fill_value = fill_value
        self.max_holes = max_holes
        self.max_spatial_size = max_spatial_size
        self.mask_only = mask_only

    def _apply_cutout(self, data, fill_val, mask_tensor=None):
        """Apply spherical cutout(s) to a tensor.

        Args:
            data: Input tensor of shape (C, D, H, W)
            fill_val: Value to fill cutout regions
            mask_tensor: Optional mask tensor for mask-only cutouts

        Returns:
            Tensor with cutout applied
        """
        result = data.clone()
        n_holes = torch.randint(self.holes, (self.max_holes or self.holes) + 1, (1,)).item()

        spatial_shape = data.shape[1:]  # (D, H, W)
        min_size = self.spatial_size if isinstance(self.spatial_size, int) else self.spatial_size[0]
        max_size = self.max_spatial_size or self.spatial_size
        max_size = max_size if isinstance(max_size, int) else max_size[0]

        for _ in range(n_holes):
            # Random size for this hole
            size = torch.randint(min_size, max_size + 1, (3,))
            radii = size.float() / 2

            # Random center (ensure hole fits in volume)
            center = [
                torch.randint(int(radii[i].item()),
                              max(spatial_shape[i] - int(radii[i].item()), int(radii[i].item()) + 1),
                              (1,)).item()
                for i in range(3)
            ]

            ellipsoid = _create_ellipsoid_mask(spatial_shape, center, radii)

            if self.mask_only and mask_tensor is not None:
                # INTERSECT with tumor mask - only affect tumor voxels
                tumor_mask = mask_tensor[0] > 0
                cutout_region = ellipsoid & tumor_mask
            else:
                cutout_region = ellipsoid

            result[:, cutout_region] = fill_val

        return result

    def apply_transform(self, subject):
        # Get mask if available for mask-only cutouts
        mask_tensor = None
        if self.mask_only and 'mask' in subject:
            mask_tensor = subject['mask'].data

            # Skip if mask is empty
            if mask_tensor is not None and not (mask_tensor > 0).any():
                return subject

        for image in self.get_images(subject):
            data = image.data
            fill_val = self.fill_value if self.fill_value is not None else float(data.min())
            result = self._apply_cutout(data, fill_val, mask_tensor)
            image.set_data(result)
        return subject

# %% ../nbs/03_vision_augment.ipynb #bc99f04d
class RandomCutout(ItemTransform):
    """Randomly erase spherical regions in 3D medical images with mask-aware placement.

    Simulates post-operative surgical cavities by filling random ellipsoid
    volumes with specified values. When mask_only=True (default), cutouts only
    affect voxels inside the segmentation mask, ensuring no healthy tissue is modified.

    Args:
        holes: Minimum number of cutout regions. Default: 1.
        max_holes: Maximum number of regions. Default: 3.
        spatial_size: Minimum cutout diameter in voxels. Default: 8.
        max_spatial_size: Maximum cutout diameter. Default: 16.
        fill: Fill value - 'min', 'mean', 'random', or float. Default: 'min'.
        mask_only: If True, cutouts only affect mask-positive voxels (tumor tissue).
            If False, cutouts can affect any voxel (original behavior). Default: True.
        p: Probability of applying transform. Default: 0.2.

    Example:
        >>> # Simulate post-op cavities only within tumor regions
        >>> tfm = RandomCutout(holes=1, max_holes=2, spatial_size=10,
        ...                    max_spatial_size=25, fill='min', mask_only=True, p=0.2)

        >>> # Original behavior - cutouts anywhere in the volume
        >>> tfm = RandomCutout(mask_only=False, p=0.2)
    """

    split_idx, order = 0, 1

    def __init__(self, holes=1, max_holes=3, spatial_size=8,
                 max_spatial_size=16, fill='min', mask_only=True, p=0.2):
        self.holes = holes
        self.max_holes = max_holes
        self.spatial_size = spatial_size
        self.max_spatial_size = max_spatial_size
        self.fill = fill
        self.mask_only = mask_only
        self.p = p

        self._tio_cutout = _TioRandomCutout(
            holes=holes, spatial_size=spatial_size,
            fill_value=None if isinstance(fill, str) else fill,
            max_holes=max_holes, max_spatial_size=max_spatial_size,
            mask_only=mask_only, p=p
        )

    @property
    def tio_transform(self):
        """Return TorchIO-compatible transform for patch-based workflows.

        Note: For mask-aware cutouts in patch workflows, the mask must be
        available in the TorchIO Subject as 'mask' key.
        """
        return self._tio_cutout

    def _get_fill_value(self, tensor):
        if self.fill == 'min': return float(tensor.min())
        elif self.fill == 'mean': return float(tensor.mean())
        elif self.fill == 'random':
            return torch.empty(1).uniform_(float(tensor.min()), float(tensor.max())).item()
        else: return self.fill

    def encodes(self, x):
        """Apply mask-aware cutout to image.

        Args:
            x: Tuple of (MedImage, target) where target is MedMask or TensorCategory

        Returns:
            Tuple of (transformed MedImage, unchanged target)
        """
        img, y_true = x

        # Probability check
        if torch.rand(1).item() > self.p:
            return img, y_true

        # Get mask data if available (as numpy for safe boolean operations)
        mask_np = None
        tumor_coords = None
        if isinstance(y_true, MedMask):
            mask_np = y_true.numpy()
            # Get coordinates of tumor voxels for mask-aware center placement
            if self.mask_only and (mask_np > 0).any():
                tumor_coords = np.argwhere(mask_np[0] > 0)  # Shape: (N, 3) for z, y, x

        # Skip cutout if mask_only=True but no mask or empty mask
        if self.mask_only:
            if mask_np is None or tumor_coords is None or len(tumor_coords) == 0:
                return img, y_true

        # Work with numpy array to avoid tensor subclass issues
        result_np = img.numpy().copy()
        spatial_shape = img.shape[1:]  # (D, H, W)
        fill_val = self._get_fill_value(img)

        n_holes = torch.randint(self.holes, self.max_holes + 1, (1,)).item()

        min_size = self.spatial_size if isinstance(self.spatial_size, int) else self.spatial_size[0]
        max_size = self.max_spatial_size or self.spatial_size
        max_size = max_size if isinstance(max_size, int) else max_size[0]

        for _ in range(n_holes):
            # Random size for this hole
            size = torch.randint(min_size, max_size + 1, (3,))
            radii = size.float() / 2

            if self.mask_only and tumor_coords is not None:
                # Pick center from within tumor region to ensure intersection
                idx = torch.randint(0, len(tumor_coords), (1,)).item()
                center = tumor_coords[idx].tolist()  # [z, y, x]
            else:
                # Random center anywhere (ensure hole fits in volume)
                center = [
                    torch.randint(int(radii[i].item()),
                                  max(spatial_shape[i] - int(radii[i].item()), int(radii[i].item()) + 1),
                                  (1,)).item()
                    for i in range(3)
                ]

            # Create ellipsoid mask (numpy)
            ellipsoid = _create_ellipsoid_mask(spatial_shape, center, radii).numpy()

            if self.mask_only and mask_np is not None:
                # INTERSECT with tumor mask - only affect tumor voxels
                tumor_mask = mask_np[0] > 0
                cutout_region = ellipsoid & tumor_mask
            else:
                cutout_region = ellipsoid

            result_np[:, cutout_region] = fill_val

        return MedImage.create(torch.from_numpy(result_np)), y_true

# %% ../nbs/03_vision_augment.ipynb #3c4f113d
class RandomElasticDeformation(CustomDictTransform):
    """Apply TorchIO `RandomElasticDeformation`."""

    def __init__(self, num_control_points=7, max_displacement=7.5,
                 image_interpolation='linear', p=0.5):
                     
        super().__init__(tio.RandomElasticDeformation(
            num_control_points=num_control_points,
            max_displacement=max_displacement,
            image_interpolation=image_interpolation,
            p=p))

# %% ../nbs/03_vision_augment.ipynb #feaecd21
class RandomAffine(CustomDictTransform):
    """Apply TorchIO `RandomAffine`."""

    def __init__(self, scales=0, degrees=10, translation=0, isotropic=False,
                 image_interpolation='linear', default_pad_value=0., p=0.5):
                     
        super().__init__(tio.RandomAffine(
            scales=scales,
            degrees=degrees,
            translation=translation,
            isotropic=isotropic,
            image_interpolation=image_interpolation,
            default_pad_value=default_pad_value,
            p=p))

# %% ../nbs/03_vision_augment.ipynb #022c90cf
class RandomFlip(CustomDictTransform):
    """Apply TorchIO `RandomFlip`."""

    def __init__(self, axes='LRAPIS', p=0.5):
        super().__init__(tio.RandomFlip(axes=axes, flip_probability=p))

# %% ../nbs/03_vision_augment.ipynb #ddd7b99b
class OneOf(CustomDictTransform):
    """Apply only one of the given transforms using TorchIO `OneOf`."""

    def __init__(self, transform_dict, p=1):
        super().__init__(tio.OneOf(transform_dict, p=p))

# %% ../nbs/03_vision_augment.ipynb #lqet5pabzy
def _compute_patch_aug_params(patch_size, target_spacing,
                              anisotropy_threshold=3.0,
                              translation_fraction=0.15):
    """Compute geometry-aware augmentation parameters from patch/spacing metadata.

    Shared logic used by both suggest_patch_augmentations (CPU/TorchIO) and
    gpu_patch_augmentations (GPU-batched). Extracts rotation degrees, translation
    offsets, and RandomAnisotropy axes from the spatial configuration.

    Anisotropy detection: if max(spacing)/min(spacing) >= threshold, rotation
    is restricted to 5 deg out-of-plane and 30 deg in-plane. Otherwise 30 deg
    symmetric.

    Args:
        patch_size: List/tuple of 3 ints -- patch dimensions.
        target_spacing: List/tuple of 3 floats -- voxel spacing.
        anisotropy_threshold: Ratio threshold for anisotropy detection (default 3.0).
        translation_fraction: Fraction of patch_size for translation (default 0.15).

    Returns:
        dict with keys:
            'degrees': tuple of 3 ints (per-axis max rotation in degrees)
            'translation': tuple of 3 ints (per-axis translation in voxels)
            'aniso_axes': tuple of ints (axes where patch_size > 1)
            'is_aniso': bool (whether spacing is anisotropic)
    """
    if len(patch_size) != 3:
        raise ValueError(f"patch_size must have 3 elements, got {len(patch_size)}")
    if len(target_spacing) != 3:
        raise ValueError(f"target_spacing must have 3 elements, got {len(target_spacing)}")

    spacing = list(target_spacing)
    ratio = max(spacing) / min(spacing)
    is_aniso = ratio >= anisotropy_threshold
    aniso_axis = spacing.index(max(spacing)) if is_aniso else None

    if is_aniso:
        degrees = [5, 5, 5]
        degrees[aniso_axis] = 30
        degrees = tuple(degrees)
    else:
        degrees = (30, 30, 30)

    translation = tuple(round(p * translation_fraction) for p in patch_size)
    aniso_axes = tuple(i for i in range(3) if patch_size[i] > 1)

    return {
        'degrees': degrees,
        'translation': translation,
        'aniso_axes': aniso_axes,
        'is_aniso': is_aniso,
    }

# %% ../nbs/03_vision_augment.ipynb #oef9rtzvbw
def _build_rotation_matrix_3d(angles_rad):
    """Build [N, 3, 3] rotation matrices from [N, 3] Euler angles (XYZ extrinsic).

    Computes R = Rz @ Ry @ Rx for each sample in the batch.

    Args:
        angles_rad: Tensor of shape [N, 3] with rotation angles in radians
                    for each axis (x, y, z).

    Returns:
        Tensor of shape [N, 3, 3] -- rotation matrices.
    """
    cos = torch.cos(angles_rad)
    sin = torch.sin(angles_rad)
    cx, cy, cz = cos[:, 0], cos[:, 1], cos[:, 2]
    sx, sy, sz = sin[:, 0], sin[:, 1], sin[:, 2]

    # R = Rz @ Ry @ Rx (combined formula)
    r00 = cy * cz;          r01 = sx * sy * cz - cx * sz;  r02 = cx * sy * cz + sx * sz
    r10 = cy * sz;          r11 = sx * sy * sz + cx * cz;  r12 = cx * sy * sz - sx * cz
    r20 = -sy;              r21 = sx * cy;                  r22 = cx * cy

    R = torch.stack([
        torch.stack([r00, r01, r02], dim=-1),
        torch.stack([r10, r11, r12], dim=-1),
        torch.stack([r20, r21, r22], dim=-1),
    ], dim=-2)
    return R

# %% ../nbs/03_vision_augment.ipynb #ts2kolv83z
class GpuPatchAugmentation:
    """GPU-batched augmentation for patch-based training.

    Operates on [B, C, D, H, W] tensors already on GPU. All operations run
    under torch.no_grad() since augmentation does not need gradient tracking.

    Transform order: spatial (affine, anisotropy, flip) then intensity
    (gamma, intensity_scale, noise, blur). Spatial transforms apply the
    same parameters to both image and mask. Intensity transforms skip the mask.

    Each transform is controlled by a parameter dict with at minimum a 'p' key
    for per-sample probability. Pass None to disable a transform.

    Args:
        affine: dict with keys 'scales', 'degrees', 'translation',
                'default_pad_value', 'p'. None to disable.
        anisotropy: dict with keys 'axes', 'downsampling', 'p'. None to disable.
        flip: dict with keys 'axes', 'p'. None to disable.
        gamma: dict with keys 'log_gamma', 'p'. None to disable.
        intensity_scale: dict with keys 'scale_range', 'p'. None to disable.
        noise: dict with keys 'std', 'p'. None to disable.
        blur: dict with keys 'std', 'p'. None to disable.

    Example::

        >>> gpu_aug = GpuPatchAugmentation(
        ...     affine={'scales': (0.7, 1.4), 'degrees': (30, 30, 30),
        ...             'translation': (25, 25, 10), 'default_pad_value': 0., 'p': 0.2},
        ...     gamma={'log_gamma': (-0.3, 0.3), 'p': 0.3},
        ...     flip={'axes': (0, 1, 2), 'p': 0.5},
        ... )
        >>> img_aug, mask_aug = gpu_aug(img_gpu, mask_gpu)
    """

    def __init__(self, affine=None, anisotropy=None, flip=None,
                 gamma=None, intensity_scale=None, noise=None, blur=None):
        self.affine = affine
        self.anisotropy = anisotropy
        self.flip = flip
        self.gamma = gamma
        self.intensity_scale = intensity_scale
        self.noise = noise
        self.blur = blur

    def __call__(self, img, mask=None):
        """Apply GPU augmentations to a batch.

        Args:
            img: Tensor [B, C, D, H, W] (float).
            mask: Tensor [B, C, D, H, W] (float), or None.

        Returns:
            Tuple (img, mask). mask is None if input was None.
        """
        with torch.no_grad():
            # Spatial transforms (same params for img and mask)
            if self.affine is not None:
                img, mask = self._apply_affine(img, mask)
            if self.anisotropy is not None:
                img, mask = self._apply_anisotropy(img, mask)
            if self.flip is not None:
                img, mask = self._apply_flip(img, mask)
            # Intensity transforms (img only)
            if self.gamma is not None:
                img = self._apply_gamma(img)
            if self.intensity_scale is not None:
                img = self._apply_intensity_scale(img)
            if self.noise is not None:
                img = self._apply_noise(img)
            if self.blur is not None:
                img = self._apply_blur(img)
        return img, mask

    def _apply_affine(self, img, mask):
        """Batched random affine via F.affine_grid + F.grid_sample."""
        cfg = self.affine
        B = img.shape[0]
        device = img.device
        dtype = img.dtype

        # Per-sample probability
        do_tfm = torch.rand(B, device=device) < cfg['p']
        if not do_tfm.any():
            return img, mask

        # Start with identity theta [B, 3, 4]
        theta = torch.zeros(B, 3, 4, device=device, dtype=dtype)
        theta[:, 0, 0] = 1.0
        theta[:, 1, 1] = 1.0
        theta[:, 2, 2] = 1.0

        idx = do_tfm.nonzero(as_tuple=True)[0]
        n = idx.shape[0]

        # Random scales per axis
        s_lo, s_hi = cfg['scales']
        scales = torch.empty(n, 3, device=device, dtype=dtype).uniform_(s_lo, s_hi)

        # Random rotation angles (degrees -> radians)
        degrees = cfg['degrees']
        if not isinstance(degrees, (list, tuple)):
            degrees = (degrees, degrees, degrees)
        angles_deg = torch.stack([
            torch.empty(n, device=device, dtype=dtype).uniform_(-degrees[0], degrees[0]),
            torch.empty(n, device=device, dtype=dtype).uniform_(-degrees[1], degrees[1]),
            torch.empty(n, device=device, dtype=dtype).uniform_(-degrees[2], degrees[2]),
        ], dim=1)  # [n, 3]
        angles_rad = angles_deg * (math.pi / 180.0)

        # Build rotation matrices [n, 3, 3]
        R = _build_rotation_matrix_3d(angles_rad)

        # Scale matrix: S @ R -> [n, 3, 3]
        S = torch.diag_embed(scales)  # [n, 3, 3]
        SR = S @ R  # [n, 3, 3]

        # Random translation (voxels -> normalized [-1, 1] coords)
        translation = cfg['translation']
        if not isinstance(translation, (list, tuple)):
            translation = (translation, translation, translation)
        spatial_size = img.shape[2:]  # (D, H, W)
        t_norm = torch.stack([
            torch.empty(n, device=device, dtype=dtype).uniform_(
                -translation[i], translation[i]
            ) * 2.0 / spatial_size[i]
            for i in range(3)
        ], dim=1)  # [n, 3]

        # Assemble theta for active samples
        theta[idx, :3, :3] = SR
        theta[idx, :3, 3] = t_norm

        # Apply grid_sample
        grid = F.affine_grid(theta, img.shape, align_corners=False)
        pad_val = cfg.get('default_pad_value', 0.)
        # For non-zero padding, shift values, sample, shift back
        if pad_val != 0.:
            img = img - pad_val
        img = F.grid_sample(img, grid, mode='bilinear',
                            padding_mode='zeros', align_corners=False)
        if pad_val != 0.:
            img = img + pad_val

        if mask is not None:
            mask = F.grid_sample(mask.float(), grid, mode='nearest',
                                 padding_mode='zeros', align_corners=False)
        return img, mask

    def _apply_anisotropy(self, img, mask):
        """Per-sample anisotropy simulation via F.interpolate.

        Downsample along a random axis with nearest interpolation,
        then upsample back with trilinear (matches TorchIO behavior).
        Only affects img, not mask (anisotropy is intensity degradation).
        """
        cfg = self.anisotropy
        B = img.shape[0]
        device = img.device
        ds_lo, ds_hi = cfg['downsampling']
        axes = cfg['axes']

        for i in range(B):
            if torch.rand(1, device=device).item() >= cfg['p']:
                continue
            # Pick random axis and downsampling factor
            axis_idx = torch.randint(len(axes), (1,), device=device).item()
            axis = axes[axis_idx]
            factor = torch.empty(1, device=device).uniform_(ds_lo, ds_hi).item()

            sample = img[i:i+1]  # [1, C, D, H, W]
            orig_size = list(sample.shape[2:])
            down_size = list(orig_size)
            down_size[axis] = max(1, round(orig_size[axis] / factor))

            # Downsample with nearest, upsample with trilinear
            down = F.interpolate(sample, size=down_size, mode='nearest')
            up = F.interpolate(down, size=orig_size, mode='trilinear',
                               align_corners=False)
            img[i:i+1] = up

        return img, mask

    def _apply_flip(self, img, mask):
        """Per-sample random flip along configured axes."""
        cfg = self.flip
        B = img.shape[0]
        device = img.device
        axes = cfg['axes']
        p = cfg['p']

        for i in range(B):
            # Each axis is independently flipped with probability p
            flip_dims = []
            for axis in axes:
                if torch.rand(1, device=device).item() < p:
                    # axis 0 -> dim 2 (D), axis 1 -> dim 3 (H), axis 2 -> dim 4 (W)
                    # but img[i] is [C, D, H, W], so axis 0 -> dim 1, etc.
                    flip_dims.append(axis + 2)  # +2 for batch and channel dims
            if flip_dims:
                img[i] = torch.flip(img[i], dims=[d - 1 for d in flip_dims])  # -1 since no batch dim
                if mask is not None:
                    mask[i] = torch.flip(mask[i], dims=[d - 1 for d in flip_dims])

        return img, mask

    def _apply_gamma(self, img):
        """Batched gamma correction with per-sample random gamma."""
        cfg = self.gamma
        B = img.shape[0]
        device = img.device
        dtype = img.dtype
        log_lo, log_hi = cfg['log_gamma']

        active = torch.rand(B, device=device) < cfg['p']
        if not active.any():
            return img

        # Only apply clamp + pow to active samples (clamp destroys negatives)
        active_idx = active.nonzero(as_tuple=True)[0]
        log_gamma = torch.empty(active_idx.shape[0], device=device, dtype=dtype).uniform_(log_lo, log_hi)
        gamma = torch.exp(log_gamma).view(-1, 1, 1, 1, 1)
        img[active_idx] = img[active_idx].clamp(min=0).pow(gamma)
        return img

    def _apply_intensity_scale(self, img):
        """Batched intensity scaling with per-sample random factors."""
        cfg = self.intensity_scale
        B = img.shape[0]
        device = img.device
        dtype = img.dtype
        s_lo, s_hi = cfg['scale_range']

        # Per-sample scale (inactive get scale=1)
        scale = torch.empty(B, device=device, dtype=dtype).uniform_(s_lo, s_hi)
        active = torch.rand(B, device=device) < cfg['p']
        scale = torch.where(active, scale, torch.ones_like(scale))

        img = img * scale.view(B, 1, 1, 1, 1)
        return img

    def _apply_noise(self, img):
        """Batched additive Gaussian noise with per-sample random std."""
        cfg = self.noise
        B = img.shape[0]
        device = img.device
        dtype = img.dtype

        std_val = cfg['std']
        if isinstance(std_val, (list, tuple)):
            std_lo, std_hi = std_val
            per_std = torch.empty(B, device=device, dtype=dtype).uniform_(std_lo, std_hi)
        else:
            per_std = torch.full((B,), std_val, device=device, dtype=dtype)

        # Zero std for inactive samples
        active = torch.rand(B, device=device) < cfg['p']
        per_std = torch.where(active, per_std, torch.zeros_like(per_std))

        noise = torch.randn_like(img) * per_std.view(B, 1, 1, 1, 1)
        img = img + noise
        return img

    def _apply_blur(self, img):
        """Batched separable 3D Gaussian blur via F.conv3d with groups trick."""
        cfg = self.blur
        B, C, D, H, W = img.shape
        device = img.device
        dtype = img.dtype

        std_val = cfg['std']
        if isinstance(std_val, (list, tuple)):
            std_lo, std_hi = std_val
        else:
            std_lo, std_hi = 0.0, std_val

        # Per-sample sigma
        sigma = torch.empty(B, device=device, dtype=dtype).uniform_(std_lo, std_hi)
        active = torch.rand(B, device=device) < cfg['p']
        if not active.any():
            return img

        # Fixed kernel size from max sigma
        max_sigma = max(std_hi, 0.01)
        kernel_radius = int(math.ceil(3 * max_sigma))
        kernel_size = 2 * kernel_radius + 1

        # Build per-sample 1D Gaussian kernels [B, kernel_size]
        x = torch.arange(-kernel_radius, kernel_radius + 1,
                          device=device, dtype=dtype)
        # Avoid division by zero for sigma=0
        safe_sigma = torch.where(active, sigma, torch.ones_like(sigma))
        kernels = torch.exp(-x.unsqueeze(0)**2 / (2 * safe_sigma.unsqueeze(1)**2))
        kernels = kernels / kernels.sum(dim=1, keepdim=True)

        # For inactive samples, use delta kernel
        delta = torch.zeros(B, kernel_size, device=device, dtype=dtype)
        delta[:, kernel_radius] = 1.0
        kernels = torch.where(active.unsqueeze(1), kernels, delta)

        # Expand kernels for all channels: [B*C, kernel_size]
        kernels_bc = kernels.unsqueeze(1).expand(B, C, kernel_size).reshape(B * C, kernel_size)

        # Reshape img for grouped convolution: [1, B*C, D, H, W]
        img_grouped = img.reshape(1, B * C, D, H, W)

        # Separable 3D convolution: D-axis, H-axis, W-axis
        pad = kernel_radius

        # D-axis: kernel shape [B*C, 1, K, 1, 1]
        k_d = kernels_bc.reshape(B * C, 1, kernel_size, 1, 1)
        img_grouped = F.pad(img_grouped, (0, 0, 0, 0, pad, pad), mode='replicate')
        img_grouped = F.conv3d(img_grouped, k_d, groups=B * C)

        # H-axis: kernel shape [B*C, 1, 1, K, 1]
        k_h = kernels_bc.reshape(B * C, 1, 1, kernel_size, 1)
        img_grouped = F.pad(img_grouped, (0, 0, pad, pad, 0, 0), mode='replicate')
        img_grouped = F.conv3d(img_grouped, k_h, groups=B * C)

        # W-axis: kernel shape [B*C, 1, 1, 1, K]
        k_w = kernels_bc.reshape(B * C, 1, 1, 1, kernel_size)
        img_grouped = F.pad(img_grouped, (pad, pad, 0, 0, 0, 0), mode='replicate')
        img_grouped = F.conv3d(img_grouped, k_w, groups=B * C)

        return img_grouped.reshape(B, C, D, H, W)

    def __repr__(self):
        parts = []
        for name in ['affine', 'anisotropy', 'flip', 'gamma',
                      'intensity_scale', 'noise', 'blur']:
            cfg = getattr(self, name)
            if cfg is not None:
                parts.append(f"{name}(p={cfg['p']})")
        return f"GpuPatchAugmentation({', '.join(parts)})"

# %% ../nbs/03_vision_augment.ipynb #pdbh1nqo0j7
def gpu_patch_augmentations(patch_size, target_spacing,
                            anisotropy_threshold=3.0,
                            translation_fraction=0.15,
                            affine_p=0.2, anisotropy_p=0.25,
                            gamma_p=0.3, intensity_scale_p=0.1,
                            noise_p=0.1, blur_p=0.2, flip_p=0.5):
    """Create GpuPatchAugmentation with nnU-Net-inspired defaults.

    Factory function that mirrors suggest_patch_augmentations but returns
    a GpuPatchAugmentation for GPU-batched operation. Uses the same shared
    parameter logic via _compute_patch_aug_params.

    Args:
        patch_size: List/tuple of 3 ints -- patch dimensions.
        target_spacing: List/tuple of 3 floats -- voxel spacing.
        anisotropy_threshold: Ratio threshold for anisotropy detection (default 3.0).
        translation_fraction: Fraction of patch_size for translation (default 0.15).
        affine_p: Probability for RandomAffine (default 0.2).
        anisotropy_p: Probability for RandomAnisotropy (default 0.25).
        gamma_p: Probability for RandomGamma (default 0.3).
        intensity_scale_p: Probability for RandomIntensityScale (default 0.1).
        noise_p: Probability for RandomNoise (default 0.1).
        blur_p: Probability for RandomBlur (default 0.2).
        flip_p: Probability for RandomFlip per axis (default 0.5).

    Returns:
        GpuPatchAugmentation instance.

    Example::

        >>> gpu_aug = gpu_patch_augmentations([128, 128, 32], [0.5, 0.5, 1.5])
        >>> dls = MedPatchDataLoaders.from_df(..., gpu_augmentation=gpu_aug)
    """
    params = _compute_patch_aug_params(
        patch_size, target_spacing, anisotropy_threshold, translation_fraction
    )

    affine_cfg = {
        'scales': (0.7, 1.4),
        'degrees': params['degrees'],
        'translation': params['translation'],
        'default_pad_value': 0.,
        'p': affine_p,
    }

    aniso_cfg = None
    if len(params['aniso_axes']) > 0:
        aniso_cfg = {
            'axes': params['aniso_axes'],
            'downsampling': (1.5, 4),
            'p': anisotropy_p,
        }

    return GpuPatchAugmentation(
        affine=affine_cfg,
        anisotropy=aniso_cfg,
        flip={'axes': (0, 1, 2), 'p': flip_p},
        gamma={'log_gamma': (-0.3, 0.3), 'p': gamma_p},
        intensity_scale={'scale_range': (0.75, 1.25), 'p': intensity_scale_p},
        noise={'std': 0.1, 'p': noise_p},
        blur={'std': (0.5, 1.0), 'p': blur_p},
    )

# %% ../nbs/03_vision_augment.ipynb #t6hak044rc
def suggest_patch_augmentations(patch_size, target_spacing,
                                anisotropy_threshold=3.0,
                                translation_fraction=0.15):
    """Suggest patch-based augmentations with nnU-Net-inspired defaults.

    Derives rotation degrees, translation, and RandomAnisotropy axes from
    patch geometry and voxel spacing. Returns a list of fastMONAI transform
    instances ready for the ``patch_tfms`` parameter in MedPatchDataLoaders.

    Anisotropy detection: if max(spacing)/min(spacing) >= threshold, rotation
    is restricted to 5 deg out-of-plane and 30 deg in-plane. Otherwise 30 deg
    symmetric. Translation is patch_size * fraction per axis.

    Args:
        patch_size: List/tuple of 3 ints -- patch dimensions.
        target_spacing: List/tuple of 3 floats -- voxel spacing.
        anisotropy_threshold: Ratio threshold for anisotropy detection (default 3.0).
        translation_fraction: Fraction of patch_size for translation (default 0.15).

    Returns:
        list: fastMONAI transform instances (7 normally, 6 if RandomAnisotropy omitted).

    Example::

        >>> patch_tfms = suggest_patch_augmentations([128, 128, 32], [0.5, 0.5, 1.5])
        >>> dls = MedPatchDataLoaders.from_df(..., patch_tfms=patch_tfms)
    """
    params = _compute_patch_aug_params(
        patch_size, target_spacing, anisotropy_threshold, translation_fraction
    )
    degrees = params['degrees']
    translation = params['translation']
    aniso_axes = params['aniso_axes']

    # For TorchIO: pass scalar 30 when isotropic (TorchIO expands to symmetric)
    if not params['is_aniso']:
        degrees = 30

    transforms = [
        RandomAffine(scales=(0.7, 1.4), degrees=degrees, translation=translation,
                     default_pad_value=0., p=0.2),
    ]

    if len(aniso_axes) > 0:
        transforms.append(RandomAnisotropy(axes=aniso_axes, downsampling=(1.5, 4), p=0.25))

    transforms.extend([
        RandomGamma(log_gamma=(-0.3, 0.3), p=0.3),
        RandomIntensityScale(scale_range=(0.75, 1.25), p=0.1),
        RandomNoise(std=0.1, p=0.1),
        RandomBlur(std=(0.5, 1.0), p=0.2),
        RandomFlip(p=0.5),
    ])

    return transforms
