# 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', 'RandomCutout',
           'RandomElasticDeformation', 'RandomAffine', 'RandomFlip', 'OneOf']

# %% ../nbs/03_vision_augment.ipynb 2
from fastai.data.all import *
from .vision_core import *
import torchio as tio
from monai.transforms import NormalizeIntensity as MonaiNormalizeIntensity

# %% ../nbs/03_vision_augment.ipynb 5
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 7
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 8
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 9
class ZNormalization(DisplayedTransform):
    """Apply TorchIO `ZNormalization`."""

    order = 0

    def __init__(self, masking_method=None, channel_wise=True):
        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 10
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 11
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
        )
    
    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 12
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 13
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 14
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):
        result = self.add_ghosts(o)
        # Handle potential complex values from k-space operations
        if result.is_complex():
            result = torch.real(result)
        return MedImage.create(result)

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

# %% ../nbs/03_vision_augment.ipynb 15
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): 
        result = self.add_spikes(o)
        # Handle potential complex values from k-space operations
        if result.is_complex():
            result = torch.real(result)
        return MedImage.create(result)
        
    def encodes(self, o: MedMask):
        return o

# %% ../nbs/03_vision_augment.ipynb 16
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 17
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 18
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 19
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 20
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

    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 21
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):
        result = self.add_motion(o)
        # Handle potential complex values from k-space operations
        if result.is_complex():
            result = torch.real(result)
        return MedImage.create(result)

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

# %% ../nbs/03_vision_augment.ipynb 22
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 23
class _TioRandomCutout(tio.IntensityTransform):
    """TorchIO-compatible RandomCutout for patch-based workflows."""

    def __init__(self, holes=1, spatial_size=8, fill_value=None,
                 max_holes=None, max_spatial_size=None, 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

    def _apply_cutout(self, data, fill_val):
        """Apply spherical cutout(s) to a tensor."""
        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)
            ]

            mask = _create_ellipsoid_mask(spatial_shape, center, radii)
            result[:, mask] = fill_val

        return result

    def apply_transform(self, 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)
            image.set_data(result)
        return subject

# %% ../nbs/03_vision_augment.ipynb 24
class RandomCutout(DisplayedTransform):
    """Randomly erase spherical regions in 3D medical images.

    Simulates post-operative surgical cavities by filling random ellipsoid
    volumes with specified values. Useful for training on pre-op images
    to generalize to post-op scans.

    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'.
        p: Probability of applying transform. Default: 0.2.

    Example:
        >>> # Simulate post-op cavities with dark spherical voids
        >>> tfm = RandomCutout(holes=1, max_holes=2, spatial_size=10,
        ...                    max_spatial_size=25, fill='min', 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', 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.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, p=p
        )

    @property
    def tio_transform(self):
        """Return TorchIO-compatible transform for patch-based workflows."""
        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, o: MedImage):
        if torch.rand(1).item() > self.p: return o
        fill_val = self._get_fill_value(o)
        result = self._tio_cutout._apply_cutout(o.clone(), fill_val)
        return MedImage.create(result)

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

# %% ../nbs/03_vision_augment.ipynb 26
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 27
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 28
class RandomFlip(CustomDictTransform):
    """Apply TorchIO `RandomFlip`."""

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

# %% ../nbs/03_vision_augment.ipynb 29
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))
