# ELEKTRONN3 - Neural Network Toolkit
#
# Copyright (c) 2017 - now
# Max Planck Institute of Neurobiology, Munich, Germany
# Authors: Martin Drawitsch

"""
Transformations (data augmentation, normalization etc.) for semantic segmantation.

Important note: The transformations here have a similar interface to
torchvision.transforms, but there are two key differences:

1. They all map (inp, target) pairs to (transformed_inp, transformed_target)
   pairs instead of just inp to transformed_inp. Most transforms don't change the target, though.
2. They exclusively operate on numpy.ndarray data instead of PIL or torch.Tensor data.

"""

from typing import Sequence, Tuple, Optional, Dict, Any, Callable, Union

import warnings
import numpy as np
import skimage.exposure
import skimage.transform

from scipy.ndimage.filters import gaussian_filter
from scipy.ndimage.interpolation import map_coordinates
from scipy.ndimage.morphology import distance_transform_edt


from elektronn3.data.transforms import random_blurring
from elektronn3.data.transforms.random import Normal, HalfNormal, RandInt

Transform = Callable[
    [np.ndarray, Optional[np.ndarray]],
    Tuple[np.ndarray, Optional[np.ndarray]]
]


class _DropSample(Exception):
    """Sample will be dropped and won't be fed into a DataLoader"""
    pass


class Identity:
    def __call__(self, inp, target=None):
        return inp, target


class Compose:
    """Composes several transforms together.

    Args:
        transforms (list of ``Transform`` objects): list of transforms to compose.


    Example:
        >>> Compose([
        >>>     Normalize(mean=(155.291411,), std=(41.812504,)),
        >>> ])
    """

    def __init__(self, transforms: Sequence[Transform]):
        self.transforms = transforms

    def __call__(self, inp, target):
        for t in self.transforms:
            inp, target = t(inp, target)
        return inp, target

    def __repr__(self):
        format_string = self.__class__.__name__ + '('
        for t in self.transforms:
            format_string += '\n    {0}'.format(t)
        format_string += '\n)'
        return format_string


class Lambda:
    """Wraps a function of the form f(x, y) = (x', y') into a transform.

    Args:
        func: A function that takes two arrays and returns a
            tuple of two arrays.

    Example:
        >>> # Multiplies inputs (x) by 255, leaves target (y) unchanged
        >>> t = Lambda(lambda x, y: (x * 255, y))

        >>> # You can also pass regular Python functions to Lambda
        >>> def flatten(x, y):
        >>>     return x.reshape(-1), y.reshape(-1)
        >>> t = Lambda(flatten)
    """
    def __init__(
            self,
            func: Callable[[np.ndarray, np.ndarray], Tuple[np.ndarray, np.ndarray]]
    ):
        self.func = func

    def __call__(
            self,
            inp: np.ndarray,
            target: np.ndarray
    ) -> Tuple[np.ndarray, np.ndarray]:
        return self.func(inp, target)


class RandomSlicewiseTransform:
    """Wraps any 2D transform and applies it per 2D slice independently in a 3D
    input-target pair.

    Works with any (..., D, H, W) memory layout given that if the ``target``
    is not ``None``, the last three dimensions of target and input tensors
    match.

    Args:
        transform: transform that works on 2D slices
        prob: Probability with which each slice is chosen to transformed by the
            specified ``transform``.

    Example::

        Here we replace each slice by zeros with p=0.1. This has an effect
        similar to the "missing section" augmentation described in
        https://arxiv.org/abs/1706.00120.

        >>> def zero_out(inp, target): return inp * 0, target
        >>> t = RandomSlicewiseTransform(zero_out, prob=0.1)
    """
    def __init__(
            self,
            transform: Callable[[np.ndarray, np.ndarray], Tuple[np.ndarray, np.ndarray]],
            prob: float = 0.1,
            inplace: bool = True
    ):
        self.transform = transform
        self.prob = prob
        assert inplace, 'Only inplace operation is supported currently'

    def __call__(
            self,
            inp: np.ndarray,
            target: Optional[np.ndarray]
    ) -> Tuple[np.ndarray, np.ndarray]:
        assert inp.ndim == 4, 'Input must be of shape (C, D, H, W)'
        if target is not None:
            assert inp.shape[-3:] == target.shape[-3:], 'Spatial shape of target must match input'
        D = inp.shape[-3]  # Extent in z dimension
        for z in range(D):
            # TODO: rand can be vectorized
            if np.random.rand() < self.prob:  # Only apply to this z slice with p=prob
                if target is None:
                    # For support of (..., D, H, W) memory layout:
                    # Indexing with [..., z, :, :] is necessary to stay agnostic
                    # to the number of pre-depth dimensions (we only know for
                    # certain that depth is the third from last dimension).
                    inp[..., z, :, :], _ = self.transform(inp[..., z, :, :], None)
                else:
                    inp[..., z, :, :], target[..., z, :, :] = self.transform(inp[..., z, :, :], target[..., z, :, :])
        return inp, target


class DropIfTooMuchBG:
    """Filter transform that skips a sample if the background class is too
    dominant in the target."""
    def __init__(self, bg_id=0, threshold=0.9, prob=1.0):
        self.bg_id = bg_id
        self.threshold = threshold
        self.prob = prob

    def __call__(
            self,
            inp: np.ndarray,
            target: np.ndarray
    ) -> Tuple[np.ndarray, np.ndarray]:
        if np.random.rand() > self.prob:
            return inp, target
        if np.sum(target == self.bg_id) / target.size > self.threshold:
            raise _DropSample
        return inp, target  # Return inp, target unmodified


class RemapTargetIDs:
    """Remap class IDs of targets to a new class mapping.
    If ``ids`` is a dict, it is used as a lookup table of the form orig_id -> changed_id.
    Each occurence of orig_id will be changed to changed_id.
    If ``ids`` is a list (deprecated), a dense remapping is performed (the given ids are
    remapped to [0, 1, 2, ..., N - 1], where N is ``len(ids)``).
    E.g. if your targets contain the class IDs [1, 3, 7, 9] but you are only
    interested in classes 1, 3 and 9 and you
    don't want to train a sparse classifier with useless outputs, you can
    use ``RemapTargetIDs([1, 3, 9])`` to translate each occurence of IDs
    [1, 3, 9] to [0, 1, 2], respectively.

    If ``reverse=True``, the mapping is inverted (useful for converting back
    to original mappings)."""

    def __init__(self, ids: Union[Sequence[int], Dict[int, int]], reverse: bool = False):
        if isinstance(ids, dict):
            ids = {int(k): int(v) for k, v in ids.items()}
        else:
            warnings.warn(
                'Passing lists to RemapTargetIDs is deprecated. Please use dicts instead (see docs).',
                DeprecationWarning
            )
        self.ids = ids
        self.reverse = reverse

    def __call__(
            self,
            inp: Optional[np.ndarray],  # returned without modifications
            target: Optional[np.ndarray]
    ) -> Tuple[Optional[np.ndarray], Optional[np.ndarray]]:
        if target is None:
            return inp, target
        remapped_target = target.copy()
        if isinstance(self.ids, dict):
            ids = self.ids if not self.reverse else {v: k for k, v in self.ids.items()}
            mask = {}
            for orig_id in ids.keys():
                mask[orig_id] = target == orig_id
            for orig_id, changed_id in ids.items():
                remapped_target[mask[orig_id]] = changed_id
        else:
            for changed_id, orig_id in enumerate(self.ids):
                if not self.reverse:
                    remapped_target[target == orig_id] = changed_id
                else:
                    remapped_target[target == changed_id] = orig_id
        return inp, remapped_target


class SmoothOneHotTarget:
    """Converts target tensors to one-hot encoding, with optional label smoothing.

    Args:
        out_channels: Number of target channels (C) in the data set.
        smooth_eps: Label smoothing strength. If ``smooth_eps=0`` (default), no
            smoothing is applied and regular one-hot tensors are returned.
            See section 7 of https://arxiv.org/abs/1512.00567
    """
    # TODO: Add example to docstring
    def __init__(self, out_channels: int, smooth_eps: float = 0.):
        assert 0 <= smooth_eps < 0.5
        self.out_channels = out_channels
        self.smooth_eps = smooth_eps

    def __call__(
            self,
            inp: np.ndarray,  # returned without modifications
            target: np.ndarray
    ) -> Tuple[np.ndarray, np.ndarray]:
        if self.smooth_eps == 0.:
            eye = np.eye(self.out_channels)
        else:
            # Create a "soft" eye where  0 is replaced by smooth_eps and 1 by (1 - smooth_eps)
            eye = np.full((self.out_channels, self.out_channels), self.smooth_eps)
            np.fill_diagonal(eye, 1. - self.smooth_eps)
        onehot = np.moveaxis(eye[target], -1, 0)
        assert np.all(onehot.argmax(0) == target)
        return inp, onehot


class DistanceTransformTarget:
    """Converts discrete binary label target tensors to their (signed)
    euclidean distance transform (EDT) representation.

    Based on the method proposed in https://arxiv.org/abs/1805.02718.

    Args:
        scale: Scalar value to divide distances before applying normalization
        normalize_fn: Function to apply to distance map for normalization.
        inverted: Invert target labels before computing transform if ``True``.
             This means the distance map will show the distance to the nearest
             foreground pixel at each background pixel location (which is the
             opposite behavior of standard distance transform).
        signed: Compute signed distance transform (SEDT), where foreground
            regions are not 0 but the negative distance to the nearest
            foreground border.
        vector: Return distance vector map instead of scalars.
    """
    def __init__(
            self,
            scale: Optional[float] = 50.,
            normalize_fn: Optional[Callable[[np.ndarray], np.ndarray]] = np.tanh,
            inverted: bool = True,
            signed: bool = True,
            vector: bool = False
    ):
        self.scale = scale
        self.normalize_fn = normalize_fn
        self.inverted = inverted
        self.signed = signed
        self.vector = vector

    def edt(self, target: np.ndarray) -> np.ndarray:
        sh = target.shape
        if target.min() == 1:  # If everything is 1, the EDT should be inf for every pixel
            nc = target.ndim if self.vector else 1
            return np.full((nc, *sh), np.inf, dtype=np.float32)

        if self.vector:
            if target.ndim == 2:
                coords = np.mgrid[:sh[0], :sh[1]]
            elif target.ndim == 3:
                coords = np.mgrid[:sh[0], :sh[1], :sh[2]]
            else:
                raise RuntimeError(f'Target shape {sh} not understood.')
            inds = distance_transform_edt(
                target, return_distances=False, return_indices=True
            ).astype(np.float32)
            dist = inds - coords
            # assert np.isclose(np.sqrt(dist[0] ** 2 + dist[1] ** 2), distance_transform_edt(target))
            return dist

        # Else: Regular scalar edt
        dist = distance_transform_edt(target).astype(np.float32)[None]
        return dist

    def __call__(
            self,
            inp: np.ndarray,  # returned without modifications
            target: Optional[np.ndarray]
    ) -> Tuple[np.ndarray, Optional[np.ndarray]]:
        if target is None:
            return inp, target
        # assert target.max() <= 1
        # Ensure np.bool dtype, invert if needed
        if self.inverted:
            target = target == 0
        else:
            target = target > 0
        dist = self.edt(target)
        if self.signed:
            # Compute same transform on the inverted target. The inverse transform can be
            #  subtracted from the original transform to get a signed distance transform.
            invdist = self.edt(~target)
            dist -= invdist
        if self.normalize_fn is not None:
            dist = self.normalize_fn(dist / self.scale)
        return inp, dist


class Normalize:
    """Normalizes inputs with supplied per-channel means and stds.

    Args:
        mean: Global mean value(s) of the inputs. Can either be a sequence
            of float values where each value corresponds to a channel
            or a single float value (only for single-channel data).
        std: Global standard deviation value(s) of the inputs. Can either
            be a sequence of float values where each value corresponds to a
            channel or a single float value (only for single-channel data).
        inplace: Apply in-place (works faster, needs less memory but overwrites
            inputs).
        channels: If ``channels`` is ``None``, the change is applied to
            all channels of the input tensor.
            If ``channels`` is a ``Sequence[int]``, change is only applied
            to the specified channels.
            E.g. with mean [a, b], std [x, y] and channels [0, 2],
            following normalizations will be allied:
            - channel 0 with mean a and std x
            - channel 2 with mean b and std y
    """
    def __init__(
            self,
            mean: Union[Sequence[float], float],
            std: Union[Sequence[float], float],
            inplace: bool = False,
            channels: Optional[Sequence[int]] = None
    ):
        self.mean = np.array(mean)
        self.std = np.array(std)
        self.inplace = inplace
        self.channels = channels
        # Unsqueeze first dimensions if mean and scalar are passed as scalars
        if self.mean.ndim == 0:
            self.mean = self.mean[None]
        if self.std.ndim == 0:
            self.std = self.std[None]

    def __call__(
            self,
            inp: np.ndarray,
            target: Optional[np.ndarray] = None  # returned without modifications
    ) -> Tuple[np.ndarray, np.ndarray]:
        if self.inplace:
            normalized = inp  # Refer to the same memory space
        else:
            normalized = inp.copy()
        channels = range(inp.shape[0]) if self.channels is None else self.channels
        if not len(channels) == self.mean.shape[0] == self.std.shape[0]:
            raise ValueError(f'mean ({self.mean.shape[0]}) and std ({self.std.shape[0]}) must have the same length as the C '
                             f'axis (number of channels) of the input ({inp.shape[0]}).')
        for c in channels:
            normalized[c] = (inp[c] - self.mean[c]) / self.std[c]
        return normalized, target

    def __repr__(self):
        return f'Normalize(mean={self.mean}, std={self.std}, inplace={self.inplace})'


# TODO: Support uniform distribution (or any distribution in general)
class RandomBrightnessContrast:
    """Randomly changes brightness and contrast of the input image.

    Args:
        brightness_std: Standard deviation of contrast change strength (bias)
        contrast_std: Standard deviation of brightness change strength (scale)
        channels: If ``channels`` is ``None``, the change is applied to
            all channels of the input tensor.
            If ``channels`` is a ``Sequence[int]``, change is only applied
            to the specified channels.
        prob: probability (between 0 and 1) with which to perform this
            augmentation. The input is returned unmodified with a probability
            of ``1 - prob``.
    """

    def __init__(
            self,
            brightness_std: float = 0.5,
            contrast_std: float = 0.5,
            channels: Optional[Sequence[int]] = None,
            prob: float = 1.0,
    ):
        if not channels:  # Support empty sequences as an alias for None
            channels = None
        self.channels = channels
        self.prob = prob
        self.brightness_gen = Normal(mean=0.0, sigma=brightness_std)
        self.contrast_gen = Normal(mean=1.0, sigma=contrast_std)

    def __call__(
            self,
            inp: np.ndarray,
            target: Optional[np.ndarray] = None  # returned without modifications
            # TODO: fast in-place version
    ) -> Tuple[np.ndarray, np.ndarray]:
        if np.random.rand() > self.prob:
            return inp, target
        channels = range(inp.shape[0]) if self.channels is None else self.channels
        augmented = inp.copy()
        for c in channels:
            a = self.contrast_gen()
            b = self.brightness_gen()
            # Formula based on tf.image.{adjust_contrast,adjust_brightness}
            # See https://www.tensorflow.org/api_docs/python/tf/image
            m = np.mean(inp[c])
            augmented[c] = a * (inp[c] - m) + m + b
        return augmented, target


# TODO: Per-channel gamma_std, infer channels from gamma_std shape? Same with AdditiveGaussianNoise
# TODO: Sample from a more suitable distribution. Due to clipping to gamma_min,
#       there is currently a strong bias towards this value.
class RandomGammaCorrection:
    """Applies random gamma correction to the input.

    Args:
        gamma_std: standard deviation of the gamma value.
        channels: If ``channels`` is ``None``, the change is applied to
            all channels of the input tensor.
            If ``channels`` is a ``Sequence[int]``, change is only applied
            to the specified channels.
        prob: probability (between 0 and 1) with which to perform this
            augmentation. The input is returned unmodified with a probability
            of ``1 - prob``.
    """
    def __init__(
            self,
            gamma_std: float = 0.5,
            gamma_min: float = 0.25,  # Prevent gamma <= 0 (0 causes zero division)
            channels: Optional[Sequence[int]] = None,
            prob: float = 1.0,
    ):
        if not channels:  # Support empty sequences as an alias for None
            channels = None
        self.channels = channels
        self.prob = prob
        self.gamma_generator = Normal(
            mean=1.0, sigma=gamma_std, bounds=(gamma_min, np.inf)
        )

    def __call__(
            self,
            inp: np.ndarray,
            target: Optional[np.ndarray] = None  # returned without modifications
            # TODO: fast in-place version
    ) -> Tuple[np.ndarray, np.ndarray]:
        if np.random.rand() > self.prob:
            return inp, target
        channels = range(inp.shape[0]) if self.channels is None else self.channels
        gcorr = inp.copy()
        for c in channels:
            gamma = self.gamma_generator()
            # adjust_gamma() requires inputs in the (0, 1) range, so the
            #  image intensity values are rescaled to (0, 1) and after
            #  applying gamma correction they are rescaled back to the original
            #  intensity range.
            orig_intensity_range = inp[c].min(), inp[c].max()
            rescaled = skimage.exposure.rescale_intensity(inp[c], out_range=(0, 1))
            gcorr01c = skimage.exposure.adjust_gamma(rescaled, gamma)  # still in (0, 1) range
            # Rescale to original (normalized) intensity range
            gcorr[c] = skimage.exposure.rescale_intensity(gcorr01c, out_range=orig_intensity_range)

        return gcorr, target


# TODO: The current necessity of intensity rescaling for normalized
#       (zero mean, unit std) inputs is really uncool. Can we circumvent this?
class RandomGrayAugment:
    r"""Performs gray value augmentations in the same way as ELEKTRONN2's
    ``greyAugment()`` function, but with additional support for inputs
    outside of the :math:`[0, 1]` intensity value range. Targets are not
    modified.

    This augmentation method randomly selects the values ``alpha``, ``beta``
    and ``gamma`` within *sensible* ranges and subsequently performs:

    - Temporarily rescaling image intensities to the :math:`[0, 1]` range
      (necessary for gamma correction).
    - Linear intensity scaling (contrast) by multiplying the input with
      :math:`\alpha = 1 + 0.3r`, where :math:`r \in \mathcal{U}[-0.5, 0.5]`.
      Value range: :math:`\alpha \in [0.85, 1.15]`.
    - Adding a constant value :math:`\beta = 0.3r`, where
      :math:`r \in \mathcal{U}[-0.5, 0.5]`.
      Value range: :math:`\beta \in [-0.15, 0.15]`.
    - Gamma correction with :math:`\gamma = 2^r`, where
      :math:`r \in \mathcal{U}[-1, 1]`.
    - Clipping all image intensity values to the range :math:`[0, 1]`.
    - Re-rescaling intensities back to the original input value range.
    """
    def __init__(
            self,
            channels: Optional[Sequence[int]] = None,
            prob: float = 1.0,
    ):
        if not channels:  # Support empty sequences as an alias for None
            channels = None
        self.channels = channels
        self.prob = prob

    def __call__(
            self,
            inp: np.ndarray,
            target: Optional[np.ndarray] = None  # returned without modifications
            # TODO: fast in-place version
    ) -> Tuple[np.ndarray, np.ndarray]:
        if np.random.rand() > self.prob:
            return inp, target

        channels = range(inp.shape[0]) if self.channels is None else self.channels
        nc = len(channels)
        aug = inp.copy()  # Copy so we don't overwrite inp
        # The calculations below have to be performed on inputs that have a
        #  value range of (0, 1), so they have to be rescaled.
        #  The augmented image will be re-rescaled to the original input value
        #  range at the end of the function.
        orig_intensity_ranges = [(inp[c].min(), inp[c].max()) for c in channels]
        for c in channels:  # TODO: Can we vectorize this?
            aug[c] = skimage.exposure.rescale_intensity(inp[c], out_range=(0, 1))

        alpha = 1 + (np.random.rand(nc) - 0.5) * 0.3  # ~ contrast
        beta = (np.random.rand(nc) - 0.5) * 0.3  # Mediates whether values are clipped for shadows or lights
        gamma = 2.0 ** (np.random.rand(nc) * 2 - 1)  # Sample from [0.5, 2]

        aug[channels] = aug[channels] * alpha[:, None, None] + beta[:, None, None]
        aug[channels] = np.clip(aug[channels], 0, 1)
        aug[channels] = aug[channels] ** gamma[:, None, None]

        for c in channels:  # Rescale to original (normalized) intensity range
            aug[c] = skimage.exposure.rescale_intensity(aug[c], out_range=orig_intensity_ranges[c])

        return aug, target


class RandomGaussianBlur:
    """Adds random gaussian blur to the input.

    Args:
        distsigma: Sigma parameter of the half-normal distribution from
            which sigmas for the gaussian blurring are drawn.
            To clear up possible confusion: The ``distsigma`` parameter does
            **not** directly parametrize the gaussian blurring, but the
            random distribution from which the blurring sigmas are drawn
            from.
        prob: probability (between 0 and 1) with which to perform this
            augmentation. The input is returned unmodified with a probability
            of ``1 - prob``.
        aniso_factor: a tuple or an array to apply the anisotropy, must
            match the dimension of the input.

    """

    def __init__(
            self,
            distsigma: float = 1,
            channels: Optional[Sequence[int]] = None,
            prob: float = 1.0,
            aniso_factor: Optional = None,
    ):
        if not channels:  # Support empty sequences as an alias for None
            channels = None
        self.channels = channels
        self.prob = prob
        self.gaussian_std = HalfNormal(sigma=distsigma)
        if aniso_factor is None or aniso_factor == 1:
            aniso_factor = np.array([1, 1, 1])
        self.aniso_factor = aniso_factor

    def __call__(
            self,
            inp: np.ndarray,
            target: Optional[np.ndarray] = None  # returned without modifications
            # TODO: fast in-place version
    ) -> Tuple[np.ndarray, np.ndarray]:
        if np.random.rand() > self.prob:
            return inp, target

        channels = range(inp.shape[0]) if self.channels is None else self.channels
        blurred_inp = inp.copy()
        for c in channels:
            self.aniso_factor = self.aniso_factor[:inp[c].ndim]
            sigma = self.gaussian_std(shape=inp[c].ndim)
            aniso_sigma = np.divide(sigma, self.aniso_factor)
            blurred_inp[c] = gaussian_filter(inp[c], sigma=aniso_sigma)

        return blurred_inp, target



class RandomBlurring:  # Warning: This operates in-place!

    _default_scheduler = random_blurring.ScalarScheduler(
        value=0.1,
        max_value=0.5,
        growth_type="lin",
        interval=500000,
        steps_per_report=1000
    )
    _default_config = {
        "probability": 0.5,
        "threshold": _default_scheduler,
        "lower_lim_region_size": [3, 6, 6],
        "upper_lim_region_size": [8, 16, 16],
        "verbose": False,
    }

    def __init__(
            self,
            config: Dict[str, Any],
            patch_shape: Optional[Sequence[int]] = None
    ):
        self.config = {**self._default_config, **config}
        if patch_shape is not None:
            random_blurring.check_random_data_blurring_config(patch_shape, **config)

    def __call__(
            self,
            inp: np.ndarray,
            target: Optional[np.ndarray] = None  # returned without modifications
    ) -> Tuple[np.ndarray, np.ndarray]:
        # In-place, overwrites inp!
        assert inp.ndim == 4, 'Currently only (C, D, H, W) inputs are supported.'
        random_blurring.apply_random_blurring(inp_sample=inp, **self.config)
        return inp, target


class AdditiveGaussianNoise:
    """Adds random gaussian noise to the input.

        Args:
            sigma: Sigma parameter of the gaussian distribution to draw from
            channels: If ``channels`` is ``None``, the noise is applied to
                all channels of the input tensor.
                If ``channels`` is a ``Sequence[int]``, noise is only applied
                to the specified channels.
            prob: probability (between 0 and 1) with which to perform this
                augmentation. The input is returned unmodified with a probability
                of ``1 - prob``.

    """

    def __init__(
            self,
            sigma: float = 0.1,
            channels: Optional[Sequence[int]] = None,
            prob: float = 1.0,
    ):
        self.channels = channels
        self.prob = prob
        self.noise_generator = Normal(mean=0, sigma=sigma)

    def __call__(
            self,
            inp: np.ndarray,
            target: Optional[np.ndarray] = None  # returned without modifications
            # TODO: fast in-place version
    ) -> Tuple[np.ndarray, np.ndarray]:
        if np.random.rand() > self.prob:
            return inp, target
        noise = np.zeros_like(inp)
        channels = range(inp.shape[0]) if self.channels is None else self.channels
        for c in channels:
            noise[c] = self.noise_generator(shape=inp[c].shape)
        noisy_inp = inp + noise
        return noisy_inp, target


class RandomCrop:
    def __init__(self, crop_shape: Sequence[int]):
        self.crop_shape = np.array(crop_shape)

    def __call__(
            self,
            inp: np.ndarray,
            target: Optional[np.ndarray] = None  # returned without modifications
    ) -> Tuple[np.ndarray, Optional[np.ndarray]]:
        ndim_spatial = len(self.crop_shape)  # Number of spatial axes E.g. 3 for (C,D,H.W)
        img_shape = np.array(inp.shape[-ndim_spatial:])
        # Number of nonspatial axes (like the C axis). Usually this is one
        ndim_nonspatial = inp.ndim - ndim_spatial
        if any(self.crop_shape > img_shape):
            raise ValueError(
                f'crop shape {self.crop_shape} can\'t be larger than image shape {img_shape}.'
            )
        # Calculate the "lower" corner coordinate of the slice
        coords_lo = np.array([
            np.random.randint(0, img_shape[i] - self.crop_shape[i] + 1)
            for i in range(ndim_spatial)
        ])
        coords_hi = coords_lo + self.crop_shape  # Upper ("high") corner coordinate.
        # Calculate necessary slice indices for reading the file
        nonspatial_slice = [  # Slicing all available content in these dims.
            slice(0, inp.shape[i]) for i in range(ndim_nonspatial)
        ]
        spatial_slice = [  # Slice only the content within the coordinate bounds
            slice(coords_lo[i], coords_hi[i]) for i in range(ndim_spatial)
        ]
        full_slice = tuple(nonspatial_slice + spatial_slice)
        inp_cropped = inp[full_slice]
        if target is None:
            return inp_cropped, target

        if target.ndim == inp.ndim - 1:  # inp: (C, [D,], H, W), target: ([D,], H, W)
            full_slice = full_slice[1:]  # Remove C axis from slice because target doesn't have it
        target_cropped = target[full_slice]
        return inp_cropped, target_cropped


def _draw_debug_grid(
        inp: np.ndarray,
        target: Optional[np.ndarray] = None,
        s: int = 16,
        v: float = 0.,
) -> Tuple[np.ndarray, Optional[np.ndarray]]:
    """Draw an ``s``-spaced grid of ``v`` values to visualize deformations."""
    if target is not None and target.ndim == inp.ndim - 1:
        target[::s] = v
    if inp.ndim == 4:
        inp[:, ::s] = v
        inp[:, :, ::s] = v
        inp[:, :, :, ::s] = v
        if target is not None:
            target[:, ::s] = v
            target[:, :, ::s] = v
            if target.ndim == 4:
                target[:, :, :, ::s] = v
    elif inp.ndim == 3:
        inp[:, ::s] = v
        inp[:, :, ::s] = v
        if target is not None:
            target[:, ::s] = v
            if target.ndim == 3:
                target[:, :, ::s] = v
    return inp, target


class ElasticTransform:
    """
    Based on https://gist.github.com/fmder/e28813c1e8721830ff9c


    Elastic deformation of images as described in Simard, 2003
    (DOI 10.1109/ICDAR.2003.1227801)

    Args:
        sigma: Sigma parameter of the gaussian smoothing performed on the
            random displacement field. High ``sigma`` values (> 4) lead
            to less randomness and more spatial consistency.
        alpha: Factor by which all random displacements are multiplied.
            Each local displacement is drawn from the range
            ``[-alpha, alpha]``, so e.g. for ``alpha=1`` you won't see
            much of an effect.
        channels: If ``channels`` is ``None``, the change is applied to
            all channels of the input tensor.
            If ``channels`` is a ``Sequence[int]``, change is only applied
            to the specified channels.
        prob: probability (between 0 and 1) with which to perform this
            augmentation. The input is returned unmodified with a probability
            of ``1 - prob``
        target_discrete_ix: list
            List of target channels that contain discrete values.
            By default (``None``), every channel is is seen as discrete (this is
            generally the case for classification tasks).
            This information is used to decide what kind of interpolation should
            be used for reading target data:

                - discrete targets are obtained by nearest-neighbor interpolation
                - non-discrete (continuous) targets are linearly interpolated.
        aniso_factor: Factor by which to divide the deformation strength in the
            z axis. E.g. if the data has half resolution in the z dimension, set
            ``aniso_factor = 2``. By default it is ``1``, so every spatial
            dimension is treated equally.
        draw_debug_grid: If ``True``, draw a 16-spaced grid into the image to
            visualize deformations. This is only for debugging purposes and
            should never be enabled during training.

    The input image should be of dimensions (C, H, W) or (C, D, H, W).
    C must be included.
    """

    def __init__(
            self,
            sigma: float = 4,
            alpha: float = 40,
            channels: Optional[Sequence[int]] = None,
            prob: float = 0.25,
            target_discrete_ix: Optional[Sequence[int]] = None,
            aniso_factor: float = 1.,
            draw_debug_grid: bool = False
    ):
        self.sigma = sigma
        self.alpha = alpha
        self.channels = channels
        self.prob = prob
        self.target_discrete_ix = target_discrete_ix
        self.aniso_factor = aniso_factor
        self.draw_debug_grid = draw_debug_grid

    def __call__(
            self,
            inp: np.ndarray,
            target: Optional[np.ndarray] = None

    ) -> Tuple[np.ndarray, Optional[np.ndarray]]:
        if np.random.rand() > self.prob:
            return inp, target

        channels = range(inp.shape[0]) if self.channels is None else self.channels

        # TODO (low priority): This could be written for n-d without explicit dimensions.
        if inp.ndim == 4:
            if self.draw_debug_grid:
                inp, target = _draw_debug_grid(inp, target)
            ish = np.array(inp.shape[-3:])
            dz = gaussian_filter(np.random.rand(*ish) * 2 - 1, self.sigma, mode="constant", cval=0) * self.alpha
            dy = gaussian_filter(np.random.rand(*ish) * 2 - 1, self.sigma, mode="constant", cval=0) * self.alpha
            dx = gaussian_filter(np.random.rand(*ish) * 2 - 1, self.sigma, mode="constant", cval=0) * self.alpha
            z, y, x = np.array(
                np.meshgrid(np.arange(ish[0]), np.arange(ish[1]), np.arange(ish[2]), indexing='ij'),
                dtype=np.float64
            )
            dz /= self.aniso_factor
            z += dz
            y += dy
            x += dx
            indices = np.reshape(z, (-1, 1)), np.reshape(y, (-1, 1)), np.reshape(x, (-1, 1))

            # If there is a target, apply the same deformation field to the target
            if target is not None:
                tsh = np.array(target.shape[-3:])
            if target is not None and np.any(ish != tsh):
                if self.draw_debug_grid:
                    inp, target = _draw_debug_grid(inp, target)
                # Crop input re-indexing arrays to the target region and transform coordinates
                #  to the target' own frame by subtracting the input-target offset
                lo = (ish - tsh) // 2
                hi = ish - lo
                tcrop = tuple([slice(lo[i], hi[i]) for i in range(3)])
                target_indices = (
                    np.reshape(z[tcrop] - lo[0], (-1, 1)),
                    np.reshape(y[tcrop] - lo[1], (-1, 1)),
                    np.reshape(x[tcrop] - lo[2], (-1, 1))
                )
            else:
                target_indices = indices
        elif inp.ndim == 3:
            ish = np.array(inp.shape[-2:])
            dy = gaussian_filter(np.random.rand(*ish) * 2 - 1, self.sigma, mode="constant", cval=0) * self.alpha
            dx = gaussian_filter(np.random.rand(*ish) * 2 - 1, self.sigma, mode="constant", cval=0) * self.alpha
            y, x = np.array(
                np.meshgrid(np.arange(ish[0]), np.arange(ish[1])),
                dtype=np.float64
            )
            y += dy
            x += dx
            indices = np.reshape(y, (-1, 1)), np.reshape(x, (-1, 1))

            # If there is a target, apply the same deformation field to the target
            if target is not None:
                tsh = np.array(target.shape[-2:])
            if target is not None and np.any(ish != tsh):
                # Crop input re-indexing arrays to the target region and transform coordinates
                #  to the target' own frame by subtracting the input-target offset
                lo = (ish - tsh) // 2
                hi = ish - lo
                tcrop = tuple([slice(lo[i], hi[i]) for i in range(2)])
                target_indices = (
                    np.reshape(y[tcrop] - lo[0], (-1, 1)),
                    np.reshape(x[tcrop] - lo[1], (-1, 1))
                )
            else:
                target_indices = indices
        else:
            raise ValueError("Input dimension not understood!")

        deformed_img = inp.copy()
        for c in channels:
            deformed_img[c] = map_coordinates(inp[c], indices, order=1).reshape(ish)

        if target is None:
            return deformed_img, target
        else:
            target_c = True  # True if the first dim of target is the number of channels
            if target.ndim == 4:  # (C, D, H, W)
                target_channels = target.shape[0]
                target_shape = target[0].shape
            elif target.ndim == 3:  # (C, H, W) or (D, H, W)
                if inp.ndim == 3:  # (C, H, W) target, (C, H, W) input
                    target_channels = target.shape[0]
                    target_shape = target[0].shape
                elif inp.ndim == 4:  # (D, H, W) target, (C, D, H, W) input
                    target_c = False
                    target_channels = 1
                    target_shape = target.shape
                else:
                    raise ValueError("Input dimension not understood!")
            elif target.ndim == 2:  # (H, W)
                target_c = False
                target_channels = 1
                target_shape = target.shape
            else:
                raise ValueError("Target dimension not understood!")

            target_discrete_ix = self.target_discrete_ix
            if target_discrete_ix is None:
                target_discrete_ix = [True for i in range(target_channels)]
            else:
                target_discrete_ix = [i in target_discrete_ix for i in range(target_channels)]

            deformed_target = target.copy()
            if target_c:
                for tc in range(target_channels):
                    target_order = 0 if target_discrete_ix[tc] is True else 1
                    deformed_target[tc] = map_coordinates(target[tc], target_indices, order=target_order).reshape(target_shape)
            else:
                target_order = 0 if target_discrete_ix[0] is True else 1
                deformed_target = map_coordinates(target, target_indices, order=target_order).reshape(target_shape)
            return deformed_img, deformed_target


class SqueezeTarget:
    """Squeeze a specified dimension in target tensors.

    (This is just needed as a workaround for the example neuro_data_cdhw data
    set, because its targets have a superfluous first dimension.)"""
    def __init__(self, dim: int):
        self.dim = dim

    def __call__(
            self,
            inp: np.ndarray,  # Returned without modifications
            target: np.ndarray,
    ) -> Tuple[np.ndarray, np.ndarray]:
        if target is None:
            return inp, target
        return inp, target.squeeze(axis=self.dim)


class RandomFlip:
    """Randomly flips spatial input and target dimensions respectively. Spatial
    dimensions are considered to occur last in the input/target shape and are
    flipped with probability p=0.5 (iid).

    Args:
        ndim_spatial: Number of spatial dimension in input, e.g.
            ``ndim_spatial=2`` for input shape (N, C, H, W)

    """
    def __init__(
            self,
            ndim_spatial: int = 2,
    ):
        self.randint = RandInt()
        self.ndim_spatial = ndim_spatial

    def __call__(
            self,
            inp: np.ndarray,
            target: Optional[np.ndarray] = None
    ) -> Tuple[np.ndarray, np.ndarray]:
        # TODO: np.flip now supports multi-dimensional flipping as of numpy 1.15
        #       So we can rewrite this with np.flip to make it more readable.
        flip_dims = self.randint(self.ndim_spatial)
        # flip all images at once
        slices_inp = tuple(
            [slice(None, None, 1) for _ in range(len(inp.shape) - self.ndim_spatial)] +
            [slice(None, None, (-1)**flip_d) for flip_d in flip_dims]
        )
        inp_flipped = inp[slices_inp].copy()
        if target is not None:
            slices_target = tuple(
                [slice(None, None, 1) for _ in range(len(target.shape) - self.ndim_spatial)] +
                [slice(None, None, (-1)**flip_d) for flip_d in flip_dims]
            )
            target_flipped = target[slices_target].copy()
        else:
            target_flipped = None
        return inp_flipped, target_flipped


# TODO: Make rotation border mode configurable
class RandomRotate2d:
    """Random rotations in the xy plane, based on scikit-image.

    If inputs are 3D images ([C,] D, H, W) Rotate them as a stack of 2D images
    by the same angle, constraining the rotation direction to the xy plane"""
    def __init__(self, angle_range=(-180, 180), prob=1):
        self.angle_range = angle_range
        self.prob = prob

    def __call__(self, inp, target):
        assert inp.ndim in (3, 4)
        if np.random.rand() > self.prob:
            return inp, target
        angle = np.random.uniform(*self.angle_range)
        rot_opts = {'angle': angle, 'preserve_range': True, 'mode': 'reflect'}

        if target is not None:
            if target.ndim == inp.ndim - 1:  # Implicit (no) channel dimension
                target_c = False
            elif target.ndim == inp.ndim:  # Explicit channel dimension
                target_c = True
            else:
                raise ValueError('Target dimension not understood.')

        def rot(inp, target):
            """Rotate in 2D space"""
            for c in range(inp.shape[0]):
                inp[c] = skimage.transform.rotate(inp[c], **rot_opts).astype(inp.dtype)[None]
            if target is None:
                return inp, target  # Return early if target shouldn't be transformed
            # Otherwise, transform target:
            if target_c:
                for c in range(target.shape[0]):
                    target[c] = skimage.transform.rotate(target[c], **rot_opts).astype(target.dtype)
            else:
                target = skimage.transform.rotate(target, **rot_opts).astype(target.dtype)
            return inp, target

        if inp.ndim == 3:  # 2D case
            rinp, rtarget = rot(inp, target)
        else:  # 3D case: Rotate each z slice separately by the same angle
            rinp = inp.copy()
            if target is not None:
                rtarget = target.copy()
                for z in range(rinp.shape[1]):
                    if target_c:
                        rinp[:, z], rtarget[:, z] = rot(inp[:, z], target[:, z])
                    else:
                        rinp[:, z], rtarget[z] = rot(inp[:, z], target[z])
            else:
                for z in range(rinp.shape[1]):
                    rinp[:, z], _ = rot(inp[:, z], None)
                rtarget = None
        return rinp, rtarget


class Clahe2d:
    def __call__(
            self,
            inp: np.ndarray,
            target: Optional[np.ndarray] = None
    ) -> Tuple[np.ndarray, np.ndarray]:
        assert inp.ndim == 3, 'Only 2D data is supported'
        orig_dtype = inp.dtype
        inp = inp.astype(np.uint8)  # equalize_adapthist() requires uint8
        clahe_inp = inp.copy()
        for c in range(inp.shape[0]):
            clahe_inp[c] = skimage.exposure.equalize_adapthist(inp[c])
        clahe_inp = clahe_inp.astype(orig_dtype)
        return clahe_inp, target


# TODO: Support other image shapes
class AlbuSeg2d:
    """Wrapper for albumentations' segmentation-compatible 2d augmentations.

    Wraps an augmentation so it can be used within elektronn3's transform pipeline.
    See https://github.com/albu/albumentations.

    If ``target`` is ``None``, it is ignored. Else, it is passed to the wrapped
    albumentations augmentation as the ``mask`` argument.

    Args:
        albu: albumentation object of type `DualTransform`.

    Example::

        >>> import albumentations
        >>> transform = AlbuSeg2d(albumentations.ShiftScaleRotate(
        ...     p=0.98, rotate_limit=180, scale_limit=0.1, interpolation=2
        ... ))
        ... # Note that interpolation=2 means cubic interpolation (-> cv2.CUBIC constant).
        ... # Don't confuse this with scipy's interpolation options.
    """

    def __init__(self, albu: 'albumentations.core.transforms_interface.DualTransform'):
        self.albu = albu

    def __call__(self, inp, target):
        if inp.ndim != 3:
            raise ValueError(f'target needs to have ndim=3, but it has shape {inp.shape}')
        inp = inp.transpose(1, 2, 0)  # (C, H, W) -> (H, W, C) because albmumentations requires it
        if target is not None:
            if target.ndim == 3 and target.shape[1:] == inp.shape[:-1]:
                target = target.transpose(1, 2, 0)  # (C, H, W) -> (H, W, C)
                augmented = self.albu(image=inp, mask=target)
                atarget = np.array(augmented['mask'], dtype=target.dtype)
                atarget = atarget.transpose(2, 0, 1)  # (H, W, C) -> (C, H, W)
            elif target.ndim == 2 and target.shape == inp.shape[:-1]:
                augmented = self.albu(image=inp, mask=target)
                atarget = np.array(augmented['mask'], dtype=target.dtype)
            else:
                raise ValueError(f'Shapes not supported. inp[:-1]: {inp.shape[:-1]}, target: {target.shape}')

        else:
            augmented = self.albu(image=inp)
            atarget = None
        ainp = np.array(augmented['image'], dtype=inp.dtype)
        ainp = ainp.transpose(2, 0, 1)  # (H, W, C) -> (C, H, W)  convert back to PyTorch layout
        return ainp, atarget


# TODO: Functional API (transforms.functional).
#       The current object-oriented interface should be rewritten as a wrapper
#       for the functional API (see implementation in torchvision).

# TODO: Extract a non-random version from each random transform.
#       E.g. RandomGammaCorrection should wrap GammaCorrection, which takes
#       the actual gamma value as an argument instead of a parametrization
#       of the random distribution from which gamma is sampled.

# TODO: Albumentations wrapper for img-to-scalar scenarios like classification
