from __future__ import annotations

from tkinter import Image
from typing import Any, Callable, Optional, Sequence, Tuple, Union

import nibabel as nib
import torch
import torchio as tio
from torch.nn import functional as F
from torchio.transforms.preprocessing.spatial.resample import Resample

from mdlu.data.modality import ImageModality

__all__ = [
    "DefaultPreprocessing",
    "ResampleOnehot",
    "ResampleAndCropOrPad",
    "CropToNonZero",
    "NNUnetNormalization",
    "RescaleIntensityPercentiles",
    "DefaultSpatialAugmentation",
    "CropOrPadPerImage",
]


class DefaultSpatialPreIntensityPreprocessingTransforms(tio.transforms.Compose):
    """The default spatial preprocessing to be running BEFORE the intensity preprocessing.

    Can optionally (by default) include a Cropping to nonzero values, a resampling to the target spacing (for masks
    either in one-hot format or using the nearest neighbor resampling).
    """

    def __init__(
        self,
        resample_onehot: bool,
        target_spacing,
        interpolation: str = "linear",
        crop_to_nonzero: str | bool | None = False,
        num_classes: int | None = None,
    ) -> None:
        """
        Args:
            resample_onehot: Whether to resample masks to the target spacing in onme-hot using the same
                interpolation scheme as for images combined with an argmax or using the nearest neighbor approach.
            target_spacing: The target spacing to resample to.
            interpolation: The interpolation scheme to use when resampling.
            crop_to_nonzero: Whether to crop to nonzero values.
            num_classes: The number of classes to use when resampling masks.

        """

        trafos = []

        if crop_to_nonzero or crop_to_nonzero is None:
            if crop_to_nonzero is not None and not isinstance(crop_to_nonzero, str):
                raise TypeError(
                    "crop_to_nonzero must be a string or None or a value evaluating to False to avoid cropping at all"
                )
            trafos.append(CropToNonZero(crop_to_nonzero))

        if resample_onehot:
            assert num_classes is not None
            trafos.append(
                ResampleOnehot(
                    num_classes=num_classes,
                    target=target_spacing,
                    image_interpolation=interpolation,
                )
            )
        else:
            trafos.append(Resample(target=target_spacing, image_interpolation=interpolation))

        super().__init__(trafos)


class DefaultSpatialPostIntensityPreprocessingTransforms(tio.transforms.Compose):
    """The default spatial preprocessing to be running AFTER the intensity preprocessing.

    Currently only consists of cropping or padding to a given target size
    """

    def __init__(self, target_size: torch.Tensor) -> None:
        """
        Args:
            target_size: The target size to crop or pad to.
        """
        trafos = [CropOrPadPerImage(target_shape=target_size)]
        super().__init__(trafos)


class DefaultIntensityPreprocessing(tio.transforms.Compose):
    """The default intensity preprocessing.

    Includes a rescaling of intensity values to certain percentiles as well as a normalization transform.
    """

    def __init__(
        self,
        target_size,
        modality: int | ImageModality | None = None,
        dataset_intensity_mean: torch.Tensor | None = None,
        dataset_intensity_std: torch.Tensor | None = None,
        norm_trafo: Callable[[tio.data.Subject], tio.data.Subject] | None = None,
    ) -> None:
        """
        Args:
            target_size: The target size to crop or pad to.
            modality: The modality to use.
            dataset_intensity_mean: The mean to use for the intensity normalization.
            dataset_intensity_std: The standard deviation to use for the intensity normalization.
            norm_trafo: A function to apply to the subject after normalization.
        """

        if modality == ImageModality.CT:
            percentiles = (0.5, 99.5)
        else:
            percentiles = (0, 100)

        if norm_trafo is None:
            norm_trafo = NNUnetNormalization(
                target_size=target_size,
                image_modality=modality,
                dataset_mean=dataset_intensity_mean,
                dataset_std=dataset_intensity_std,
            )

        trafos = [RescaleIntensityPercentiles(percentiles=percentiles), norm_trafo]
        super().__init__(trafos)


class DefaultPreprocessing(tio.transforms.Compose):
    """The default preprocessing. Consists of the following tranforms (in Order):

    - Transform all images and label images to a canonical RAS+ orientation
    - Copy Affine from image to others (like labels). There can be inaccuracies using different loaders or annotation tools
        which could cause errors downstream if not unifying now.
    - Cropping to nonzero values (optional)
    - Resampling to the target spacing (for masks: either one-hot or nearest neighbor)
    - Rescaling of intensity values to certain percentiles
    - Normalization of intensity values
    - Cropping or padding to target size
    """

    def __init__(
        self,
        target_spacing: Sequence[float],
        target_size: Sequence[int],
        modality: int | ImageModality | None = None,
        dataset_intensity_mean: torch.Tensor | None = None,
        dataset_intensity_std: torch.Tensor | None = None,
        interpolation: str = "linear",
        affine_key: None | str = "data",
        num_classes: int | None = None,
    ):
        """
        Args:
            target_spacing: The target spacing to resample to.
            target_size: The target size to crop or pad to.
            modality: The modality to use.
            dataset_intensity_mean: The mean to use for the intensity normalization.
            dataset_intensity_std: The standard deviation to use for the intensity normalization.
            interpolation: The interpolation scheme to use when resampling.
            affine_key: The key to use for copying the affine from.
            num_classes: The number of classes to use when resampling masks.

        """
        trafos = [
            tio.transforms.ToCanonical(),
        ]

        if affine_key is not None:
            trafos.append(tio.transforms.preprocessing.CopyAffine(target=affine_key))

        trafos += [
            DefaultSpatialPreIntensityPreprocessingTransforms(
                resample_onehot=True,
                num_classes=num_classes,
                target_spacing=target_spacing,
                interpolation=interpolation,
            ),
            DefaultIntensityPreprocessing(
                modality=modality,
                target_size=target_size,
                dataset_intensity_mean=dataset_intensity_mean,
                dataset_intensity_std=dataset_intensity_std,
            ),
            DefaultSpatialPostIntensityPreprocessingTransforms(target_size=target_size),
        ]

        return super().__init__(trafos)


class ResampleOnehot(Resample):
    """Resample a mask to the target spacing using one-hot encoding and same interpolation scheme as for images.

    .. note::
        For additional arguments have a look at :class:`torchio.transforms.preprocessing.spatial.resample.Resample`.
    """

    def __init__(self, num_classes: int, *args: Any, **kwargs: Any) -> None:
        """
        Args:
            num_classes: The number of classes to use when resampling masks.
        """
        super().__init__(*args, **kwargs)
        self.to_onehot_trafo = tio.transforms.OneHot(num_classes=num_classes)
        self.to_categorical = LabelToCategorical()

    def apply_transform(self, subject: tio.data.Subject) -> tio.data.Subject:
        """Applies the resampling to the whole subject.

        Args:
            subject: The subject to apply the transform to.

        Returns:
            The transformed subject.
        """
        subject = self.to_onehot_trafo(subject)
        temp_scalar_images = []

        for k, v in subject.get_images_dict(intensity_only=False).items():
            if not isinstance(v, tio.data.LabelMap):
                continue
            subject[k] = tio.data.ScalarImage(tensor=v.data, affine=v.affine)
            temp_scalar_images.append(k)
        subject = super().apply_transform(subject)
        for k in temp_scalar_images:
            subject[k] = tio.data.LabelMap(tensor=subject[k].data, affine=subject[k].affine)

        subject = self.to_categorical(subject)
        return subject


class ResampleAndCropOrPad(tio.transforms.Transform):
    """Resampling to a custom target spacing and cropping or padding to a custom target size.

    Handles each image individually since depending on initial shape and spacing the in-between results may differ.
    """

    def __init__(
        self,
        target_spacing: tuple[float, float, float],
        target_size: tuple[int, int, int],
        num_classes: int,
        interpolation: str = "linear",
        **kwargs,
    ) -> None:
        """
        Args:
            target_spacing: The target spacing to resample to.
            target_size: The target size to crop or pad to.
            num_classes: The number of classes to use when resampling masks.
            interpolation: The interpolation scheme to use when resampling.
        """
        super().__init__(**kwargs)

        self._resample = ResampleOnehot(
            num_classes,
            target_spacing,
            interpolation=interpolation,
            p=self.probability,
            copy=self.copy,
            include=self.include,
            exclude=self.exclude,
            keys=None,
            keep=self.keep,
        )
        self._crop_or_pad = tio.transforms.CropOrPad(
            target_size,
            p=self.probability,
            copy=self.copy,
            include=self.include,
            exclude=self.exclude,
            keys=None,
            keep=self.keep,
        )

    def apply_transform(self, subject: tio.data.Subject) -> tio.data.Subject:
        """Applies the resampling and cropping or padding to the whole subject.

        Args:
            subject: The subject to apply the transform to.

        Returns:
            The transformed subject.
        """

        # handle each image in subject individually to compensate
        # for misalignment of images in checks inbetween the steps!
        for k, v in subject.get_images_dict(intensity_only=False).items():
            curr_sub = tio.data.Subject(k=v)
            curr_sub = self._resample(curr_sub)
            curr_sub = self._crop_or_pad(curr_sub)

            subject[k] = curr_sub[k]
        return subject


class CropToNonZero(tio.transforms.Transform):
    """Crops the image to a bounding box containing all nonzero voxels."""

    def __init__(self, key: str | None = None, **kwargs: Any) -> None:
        """
        Args:
            key: The key specifying the image to use.
        """
        super().__init__(**kwargs)
        self.key = key

    @staticmethod
    def crop_to_nonzero(image: tio.data.Image, *additional_images: tio.data.Image, **padding_kwargs) -> None:
        """crops tensor to non-zero and additional tensor to the same part of the tensor if given.

        Args:
            tensor: the tensor to crop
            additional_tensor: an additional tensor to crop to the same region
                as :attr`tensor`
            **padding_kwargs: keyword arguments to controll the necessary padding
        """

        centers, ranges = extract_nonzero_bounding_box_from_tensor(image.tensor[None])

        low = (centers - ranges / 2)[0].cpu().detach().numpy()

        for img in (image, *additional_images):
            result = crop_tensor_to_bounding_box(
                img.tensor[None],
                centers=centers,
                ranges=ranges,
                **padding_kwargs,
            )

            img.set_data(result[0])

            img_affine = img.affine.copy()
            new_origin = nib.affines.apply_affine(img.affine, low)
            img_affine[:3, 3] = new_origin
            img.affine = img_affine

    def apply_transform(self, subject: tio.data.Subject) -> tio.data.Subject:
        """Applies the transform to the whole subject.

        Args:
            subject: The subject to apply the transform to.

        Returns:
            The transformed subject.
        """
        intensity_images = subject.get_images_dict(intensity_only=True)
        all_images = subject.get_images_dict(intensity_only=False)

        if self.key is None:
            images = intensity_images
        else:
            images = {self.key: subject[self.key]}

        assert (
            len(images) == 1
        ), "Multiple images found, this may lead to inconsistent cropping behavior. Pleas specify a reference image key"
        # TODO: Add reference image key to default trafos and dataset

        seg_images = {k: v for k, v in all_images.items() if k not in intensity_images}
        for v in images.values():
            self.crop_to_nonzero(v, *seg_images.values())

        return subject


class NNUnetNormalization(tio.transforms.ZNormalization):
    """Normalization similar to the one prposed in the nnUNet paper.

    - If the image is not a CT and not more than 25% of the image are cropped away,
        use the usual Z-Normalization across the whole image.
    - If the image is not a CT and more than 25% of the iamge are cropped away,
        extract the mean and standard deviation of from the non-cropped parts of the image only and
        perform Z-Normalization using these values.
    - If the image is a CT, use the passed mean and standard deviation values computed over the whole dataset
        and use those for Z-Normalization.
    """

    def __init__(
        self,
        target_size: Sequence[int],
        image_modality: int | ImageModality | None = None,
        dataset_mean: torch.Tensor | None = None,
        dataset_std: torch.Tensor | None = None,
    ):
        """
        Args:
            target_size: The target size to crop or pad to.
            image_modality: The image modality to use.
            dataset_mean: The mean to use for the dataset.
            dataset_std: The standard deviation to use for the dataset.
        """
        super().__init__()

        if isinstance(image_modality, int):
            image_modality = ImageModality(image_modality)
        self.image_modality = image_modality
        self.dataset_mean = dataset_mean
        self.dataset_std = dataset_std
        self.num_target_elems = torch.tensor(target_size).prod()
        self.crop_or_pad_trafo = CropOrPadPerImage(target_shape=target_size)

    def apply_normalization(self, subject: tio.data.Subject, image_name: str, mask: torch.Tensor) -> None:
        """Applies the normalization to the image.

        Args:
            subject: The subject to apply the normalization to.
            image_name: The name of the image to apply the normalization to.
            mask: The mask to use for the normalization.
        """

        # no image modality is given, use default (non-CT) behavior
        if self.image_modality is None or self.image_modality != ImageModality.CT:
            # not more than 25% of the voxels are cropped away
            if torch.tensor(subject.spatial_shape).prod() * 0.75 < self.num_target_elems:
                return super().apply_normalization(subject, image_name, mask)

            cropped_tensor = self.crop_or_pad_trafo(
                tio.data.Subject({image_name: tio.data.ScalarImage(tensor=subject[image_name].data)})
            )[image_name].data
            # masking for foreground
            mask = cropped_tensor > 0
            values = cropped_tensor[mask].float()
            mean = values.mean()
            std = values.std()

        else:
            if self.dataset_mean is None or self.dataset_std is None:
                raise ValueError("dataset mean and std can only be None if the dataset modality is not CT.")

            mean = self.dataset_mean
            std = self.dataset_std

        if std == 0:
            mean = self.dataset_mean
            std = self.dataset_std

        image = subject[image_name]
        tensor = image.data.to(mean)
        tensor -= mean
        tensor /= std
        image.set_data(tensor)


class LabelToCategorical(tio.transforms.preprocessing.label.label_transform.LabelTransform):
    """Converts a One-Hot encoded label to a categorical label."""

    def apply_transform(self, subject: tio.data.Subject) -> tio.data.Subject:
        """Applies the transform to the whole subject.

        Args:
            subject: The subject to apply the transform to.

        Returns:
            The transformed subject.
        """

        for k, v in self.get_images_dict(subject).items():
            if isinstance(v, tio.data.LabelMap):
                v.set_data(v.data.argmax(0).unsqueeze(0))
        return subject


class RescaleIntensityPercentiles(tio.transforms.RescaleIntensity):
    """Rescales the intensity of the image to the given percentiles."""

    def rescale(self, tensor: torch.Tensor, mask: torch.Tensor, image_name: str) -> torch.Tensor:
        """Rescales the intensity of the image to the given percentiles.

        Args:
            tensor: The tensor to rescale.
            mask: The mask to use for the rescaling.
            image_name: The name of the image to rescale.
        """
        if self.percentiles == (0, 100):
            return tensor
        else:
            self.out_max = tensor.max().item()
            self.out_min = tensor.min().item()
            return super().rescale(tensor, mask, image_name)


class DefaultSpatialAugmentation(tio.transforms.Compose):
    """Default spatial augmentation.

    - Approximates a random anisotropy
    - Randomly flips the image
    - Applies a random affine transformation
    - Optionally includes a random elastic deformation
    """

    def __init__(self, image_modality: ImageModality | int, include_deformation: bool = False, p: float = 0.25) -> None:
        """
        Args:
            image_modality: The image modality to use.
            include_deformation: Whether to include elastic deformation.
            p: The probability of applying every single transformation.
        """
        trafos = []
        if image_modality is not None:
            trafos += [
                tio.transforms.RandomAnisotropy(
                    axes=tuple(range(ImageModality.get_dimensionality(image_modality))),
                    p=p,
                ),
                tio.transforms.RandomFlip(tuple(range(ImageModality.get_dimensionality(image_modality))), p=p),
            ]

        trafos.append(tio.transforms.RandomAffine(p=p))

        if include_deformation:
            trafos.append(tio.transforms.RandomElasticDeformation(p=p))

        super().__init__(trafos)


class DefaultIntensityAugmentation(tio.transforms.Compose):
    """Default intensity augmentation.

    - Random Ghosting Artifacts (for MRI only)
    - Random Spiking Artifacts (for MRI only)
    - Random Motion Artifacts
    - Random Bias Fields
    - Random Gamma Transform
    - Random Blurring
    - Random Noise
    - Random Swapping of Patches
    """

    def __init__(self, image_modality: ImageModality | int, p: float = 0.25) -> None:
        """
        Args:
            image_modality: The image modality to use.
            p: The probability of applying every single transformation.

        """
        trafos = []

        if image_modality is not None and image_modality == ImageModality.MR:
            trafos.append(tio.transforms.RandomGhosting(p=p))
            trafos.append(tio.transforms.RandomSpike(p=p))

        trafos += [
            tio.transforms.RandomMotion(p=p),
            tio.transforms.RandomBiasField(p=p),
            tio.transforms.RandomGamma(p=p),
            tio.transforms.RandomBlur(p=p),
            tio.transforms.RandomNoise(p=p),
            tio.transforms.RandomSwap(p=p),
        ]
        super().__init__(trafos)


class DefaultAugmentation(tio.transforms.Compose):
    """Default augmentation.

    - Approximates a random anisotropy
    - Randomly flips the image
    - Applies a random affine transformation
    - Optionally includes a random elastic deformatio
    - Random Ghosting Artifacts (for MRI only)
    - Random Spiking Artifacts (for MRI only)
    - Random Motion Artifacts
    - Random Bias Fields
    - Random Gamma Transform
    - Random Blurring
    - Random Noise
    - Random Swapping of Patches
    """

    def __init__(
        self,
        image_modality: ImageModality | int,
        include_deformation: bool = False,
        spatial_prob: float = 0.25,
        intensity_prob: float = 0.5,
    ) -> None:
        """
        Args:
            image_modality: The image modality to use.
            include_deformation: Whether to include elastic deformation.
            spatial_prob: The probability of applying every single spatial transformation.
            intensity_prob: The probability of applying every single intensity transformation.
        """
        trafos = [
            DefaultSpatialAugmentation(
                image_modality=image_modality,
                include_deformation=include_deformation,
                p=spatial_prob,
            ),
            DefaultIntensityAugmentation(image_modality=image_modality, p=intensity_prob),
        ]
        super().__init__(trafos)


class Transform3DTo2D(tio.transforms.Transform):
    """Transforms a 3D image to a 2D image with the number of channels equal to the depth beforehand."""

    def apply_transform(self, subject: tio.data.Subject) -> tio.data.Subject:
        """Applies the transform to the whole subject.

        Args:
            subject: The subject to apply the transform to.

        Returns:
            The transformed subject.
        """
        for k, v in subject.get_images_dict.items():
            tensor = v.data
            orig_shape = v.data.shape
            tensor = tensor.view(orig_shape[3] * orig_shape[0], orig_shape[1], orig_shape[2], 1)
            v.set_data(tensor)
            subject[f"_{k}_orig_shape"] = orig_shape
        return subject


class Transform2DBackTo3D(tio.transforms.Transform):
    """Transforms a 2D image back to a 3D image with the original shape restored."""

    def apply_transform(self, subject: tio.data.Subject) -> tio.data.Subject:
        """Applies the transform to the whole subject.

        Args:
            subject: The subject to apply the transform to.

        Returns:
            The transformed subject.
        """
        for k, v in subject.get_images_dict.items():
            tensor = v.data
            orig_shape = subject[f"_{k}_orig_shape"]
            tensor = tensor.view(orig_shape)
            v.set_data(tensor)
            subject.pop(f"_{k}_orig_shape")
        return subject


def extract_nonzero_bounding_box_from_tensor(
    tensor: torch.Tensor,
) -> tuple[torch.Tensor, torch.Tensor]:
    """Extracts the bounding box of nonzero elements from a tensor.

    Args:
        tensor: The tensor to extract the bounding box from.

    Returns:
        The bounding box of nonzero elements (as centers and ranges).
    """

    # create empty ranges
    ranges = torch.zeros(
        tensor.size(0),
        tensor.ndim - 2,
        dtype=torch.long,
        device=tensor.device,
    )
    # create empty centers
    centers = torch.zeros(
        tensor.size(0),
        tensor.ndim - 2,
        dtype=torch.long,
        device=tensor.device,
    )
    for idx, _tensor in enumerate(tensor):
        bounds = torch.tensor(
            [(tmp.min(), tmp.max()) for tmp in _tensor.nonzero(as_tuple=True)],
            device=tensor.device,
        )
        # ensure this is an even number
        ranges[idx, :] = (torch.ceil((bounds[:, 1] - bounds[:, 0]).float() / 2) * 2).long()[
            1:
        ]  # extract max range of all channels
        centers[idx, :] = bounds[1:, 0] + ranges[idx] // 2

    return centers, ranges


def crop_tensor_to_bounding_box(
    tensor: torch.Tensor,
    centers: torch.Tensor,
    ranges: torch.Tensor,
    additional_tensor: torch.Tensor | None = None,
    **padding_kwargs,
) -> tuple[torch.Tensor, torch.Tensor] | torch.Tensor:
    """Crops a tensor to a bounding box.

    Args:
        tensor: The tensor to crop.
        centers: The centers of the bounding box.
        ranges: The ranges of the bounding box.
        additional_tensor: An additional tensor to crop with the same bounbding box.
        **padding_kwargs: Additional arguments to pass to the padding function.

    Returns:
        The cropped tensor optionally (if passed) along with the cropped additional tensor.
    """
    max_range = ranges.max(0)[0]
    out = torch.zeros(
        tensor.size(0),
        tensor.size(1),
        *max_range.tolist(),
        device=tensor.device,
        dtype=tensor.dtype,
    )

    if additional_tensor is not None:
        additional_out = torch.zeros(
            additional_tensor.size(0),
            additional_tensor.size(1),
            *max_range.tolist(),
            device=additional_tensor.device,
            dtype=additional_tensor.dtype,
        )

    # calciulate the absolute minimum range
    abs_min_range = (centers - max_range / 2.0).min(0)[0]

    # calculate padding on the first side per dim
    first_paddings = torch.where(abs_min_range < 0, abs_min_range.abs(), torch.zeros_like(abs_min_range))

    # calculate padding on the last side per dim
    last_paddings = torch.clamp(
        (centers + max_range / 2.0).max(0)[0]
        - torch.tensor(tensor.shape[2:], device=centers.device, dtype=centers.dtype),
        0,
    )

    # combine paddings
    total_paddings = []

    for idx in range(len(first_paddings)):
        total_paddings += [first_paddings[idx], last_paddings[idx]]

    total_paddings = torch.stack(total_paddings).long().tolist()

    # add padding offset to centers
    centers = centers + first_paddings[None]

    # pad tensor
    tensor = F.pad(tensor, total_paddings, **padding_kwargs)

    # pad additional tensor if given
    if additional_tensor is not None:
        additional_tensor = F.pad(additional_tensor, total_paddings, **padding_kwargs)

    # re-assign tensors
    for idx in range(out.size(0)):
        mins = (centers[idx] - max_range / 2.0).long()
        maxs = (centers[idx] + max_range / 2.0).long()
        if len(mins) == 2:
            out[idx] = tensor[idx][..., mins[0] : maxs[0], mins[1] : maxs[1]]
            if additional_tensor is not None:
                additional_out[idx] = additional_tensor[idx][..., mins[0] : maxs[0], mins[1] : maxs[1]]
        elif len(mins) == 3:
            out[idx] = tensor[idx][..., mins[0] : maxs[0], mins[1] : maxs[1], mins[2] : maxs[2]]
            if additional_tensor is not None:
                additional_out[idx] = additional_tensor[idx][
                    ..., mins[0] : maxs[0], mins[1] : maxs[1], mins[2] : maxs[2]
                ]
    # remove channel dim for mask_out if mask had no channel dim
    if additional_tensor is not None and additional_tensor.ndim == tensor.ndim - 1:
        additional_out = additional_out[:, 0]

    if additional_tensor is None:
        return out
    return out, additional_out


def crop_tensor_to_nonzero(
    tensor: torch.Tensor,
    additional_tensor: torch.Tensor | None = None,
    **padding_kwargs,
) -> tuple[torch.Tensor, torch.Tensor] | torch.Tensor:
    """
    crops tensor to non-zero and additional tensor to the same part of the
    tensor if given
    Args:
        tensor: the tensor to crop
        additional_tensor: an additional tensor to crop to the same region
            as :attr`tensor`
        **padding_kwargs: keyword arguments to controll the necessary padding
    Returns:
        torch.Tensor: the cropped tensor
        Optional[torch.Tensor]: the cropped additional tensor,
            only returned if passed
    """
    centers, ranges = extract_nonzero_bounding_box_from_tensor(tensor)

    return crop_tensor_to_bounding_box(
        tensor,
        centers=centers,
        ranges=ranges,
        additional_tensor=additional_tensor,
        **padding_kwargs,
    )


class CropOrPadPerImage(tio.transforms.CropOrPad):
    """Crops or pads every single image individually to the given shape."""

    def apply_transform(self, subject: tio.data.Subject) -> tio.data.Subject:
        """
        Args:
            subject: the subject to apply the transform to

        Returns:
            the subject with the transformed data

        """
        for k, v in subject.get_images_dict(intensity_only=False, include=self.include, exclude=self.exclude).items():
            part_sub = type(subject)({k: v})
            subject[k] = super().apply_transform(part_sub)[k]
        return subject


# TODO: Add Transform to cut only one side of the image and flip it.
