"""
This file regroups several custom keras layers used in the generation model:
    - RandomSpatialDeformation,
    - RandomCrop,
    - RandomFlip,
    - SampleConditionalGMM,
    - SampleResolution,
    - GaussianBlur,
    - DynamicGaussianBlur,
    - MimicAcquisition,
    - BiasFieldCorruption,
    - IntensityAugmentation,
    - DiceLoss,
    - WeightedL2Loss,
    - ResetValuesToZero,
    - ConvertLabels,
    - PadAroundCentre,
    - MaskEdges
    - ImageGradients
    - RandomDilationErosion


If you use this code, please cite the first SynthSeg paper:
https://github.com/BBillot/lab2im/blob/master/bibtex.bib

Copyright 2020 Benjamin Billot

Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in
compliance with the License. You may obtain a copy of the License at
https://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software distributed under the License is
distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
implied. See the License for the specific language governing permissions and limitations under the
License.
"""

# python imports
import keras
import numpy as np
import tensorflow as tf
import tensorflow.keras.backend as K
from keras.layers import Layer

# project imports
from lamareg.ext.lab2im import utils
from lamareg.ext.lab2im import edit_tensors as l2i_et

# third-party imports
from lamareg.ext.neuron import utils as nrn_utils
import lamareg.ext.neuron.layers as nrn_layers


class RandomSpatialDeformation(Layer):
    """This layer spatially deforms one or several tensors with a combination of affine and elastic transformations.
    The input tensors are expected to have the same shape [batchsize, shape_dim1, ..., shape_dimn, channel].
    The non-linear deformation is obtained by:
    1) a small-size SVF is sampled from a centred normal distribution of random standard deviation.
    2) it is resized with trilinear interpolation to half the shape of the input tensor
    3) it is integrated to obtain a diffeomorphic transformation
    4) finally, it is resized (again with trilinear interpolation) to full image size
    :param scaling_bounds: (optional) range of the random scaling to apply. The scaling factor for each dimension is
    sampled from a uniform distribution of predefined bounds. Can either be:
    1) a number, in which case the scaling factor is independently sampled from the uniform distribution of bounds
    [1-scaling_bounds, 1+scaling_bounds] for each dimension.
    2) a sequence, in which case the scaling factor is sampled from the uniform distribution of bounds
    (1-scaling_bounds[i], 1+scaling_bounds[i]) for the i-th dimension.
    3) a numpy array of shape (2, n_dims), in which case the scaling factor is sampled from the uniform distribution
     of bounds (scaling_bounds[0, i], scaling_bounds[1, i]) for the i-th dimension.
    4) False, in which case scaling is completely turned off.
    Default is scaling_bounds = 0.15 (case 1)
    :param rotation_bounds: (optional) same as scaling bounds but for the rotation angle, except that for cases 1
    and 2, the bounds are centred on 0 rather than 1, i.e. [0+rotation_bounds[i], 0-rotation_bounds[i]].
    Default is rotation_bounds = 15.
    :param shearing_bounds: (optional) same as scaling bounds. Default is shearing_bounds = 0.012.
    :param translation_bounds: (optional) same as scaling bounds. Default is translation_bounds = False, but we
    encourage using it when cropping is deactivated (i.e. when output_shape=None in BrainGenerator).
    :param enable_90_rotations: (optional) whether to rotate the input by a random angle chosen in {0, 90, 180, 270}.
    This is done regardless of the value of rotation_bounds. If true, a different value is sampled for each dimension.
    :param nonlin_std: (optional) maximum value of the standard deviation of the normal distribution from which we
    sample the small-size SVF. Set to 0 if you wish to completely turn the elastic deformation off.
    :param nonlin_scale: (optional) if nonlin_std is not False, factor between the shapes of the input tensor
    and the shape of the input non-linear tensor.
    :param inter_method: (optional) interpolation method when deforming the input tensor. Can be 'linear', or 'nearest'
    :param prob_deform: (optional) probability to apply spatial deformation
    """

    def __init__(
        self,
        scaling_bounds=0.15,
        rotation_bounds=10,
        shearing_bounds=0.02,
        translation_bounds=False,
        enable_90_rotations=False,
        nonlin_std=4.0,
        nonlin_scale=0.0625,
        inter_method="linear",
        prob_deform=1,
        **kwargs,
    ):

        # shape attributes
        self.n_inputs = 1
        self.inshape = None
        self.n_dims = None
        self.small_shape = None

        # deformation attributes
        self.scaling_bounds = scaling_bounds
        self.rotation_bounds = rotation_bounds
        self.shearing_bounds = shearing_bounds
        self.translation_bounds = translation_bounds
        self.enable_90_rotations = enable_90_rotations
        self.nonlin_std = nonlin_std
        self.nonlin_scale = nonlin_scale

        # boolean attributes
        self.apply_affine_trans = (
            (self.scaling_bounds is not False)
            | (self.rotation_bounds is not False)
            | (self.shearing_bounds is not False)
            | (self.translation_bounds is not False)
            | self.enable_90_rotations
        )
        self.apply_elastic_trans = self.nonlin_std > 0
        self.prob_deform = prob_deform

        # interpolation methods
        self.inter_method = inter_method

        super(RandomSpatialDeformation, self).__init__(**kwargs)

    def get_config(self):
        config = super().get_config()
        config["scaling_bounds"] = self.scaling_bounds
        config["rotation_bounds"] = self.rotation_bounds
        config["shearing_bounds"] = self.shearing_bounds
        config["translation_bounds"] = self.translation_bounds
        config["enable_90_rotations"] = self.enable_90_rotations
        config["nonlin_std"] = self.nonlin_std
        config["nonlin_scale"] = self.nonlin_scale
        config["inter_method"] = self.inter_method
        config["prob_deform"] = self.prob_deform
        return config

    def build(self, input_shape):

        if not isinstance(input_shape, list):
            inputshape = [input_shape]
        else:
            self.n_inputs = len(input_shape)
            inputshape = input_shape
        self.inshape = inputshape[0][1:]
        self.n_dims = len(self.inshape) - 1

        if self.apply_elastic_trans:
            self.small_shape = utils.get_resample_shape(
                self.inshape[: self.n_dims], self.nonlin_scale, self.n_dims
            )
        else:
            self.small_shape = None

        self.inter_method = utils.reformat_to_list(
            self.inter_method, length=self.n_inputs, dtype="str"
        )

        self.built = True
        super(RandomSpatialDeformation, self).build(input_shape)

    def call(self, inputs, **kwargs):

        # reformat inputs and get its shape
        if self.n_inputs < 2:
            inputs = [inputs]
        types = [v.dtype for v in inputs]
        inputs = [tf.cast(v, dtype="float32") for v in inputs]
        batchsize = tf.split(tf.shape(inputs[0]), [1, self.n_dims + 1])[0]

        # initialise list of transforms to operate
        list_trans = list()

        # add affine deformation to inputs list
        if self.apply_affine_trans:
            affine_trans = utils.sample_affine_transform(
                batchsize,
                self.n_dims,
                self.rotation_bounds,
                self.scaling_bounds,
                self.shearing_bounds,
                self.translation_bounds,
                self.enable_90_rotations,
            )
            list_trans.append(affine_trans)

        # prepare non-linear deformation field and add it to inputs list
        if self.apply_elastic_trans:

            # sample small field from normal distribution of specified std dev
            trans_shape = tf.concat(
                [batchsize, tf.convert_to_tensor(self.small_shape, dtype="int32")],
                axis=0,
            )
            trans_std = tf.random.uniform((1, 1), maxval=self.nonlin_std)
            elastic_trans = tf.random.normal(trans_shape, stddev=trans_std)

            # reshape this field to half size (for smoother SVF), integrate it, and reshape to full image size
            resize_shape = [
                max(int(self.inshape[i] / 2), self.small_shape[i])
                for i in range(self.n_dims)
            ]
            elastic_trans = nrn_layers.Resize(
                size=resize_shape, interp_method="linear"
            )(elastic_trans)
            elastic_trans = nrn_layers.VecInt()(elastic_trans)
            elastic_trans = nrn_layers.Resize(
                size=self.inshape[: self.n_dims], interp_method="linear"
            )(elastic_trans)
            list_trans.append(elastic_trans)

        # apply deformations and return tensors with correct dtype
        if self.apply_affine_trans | self.apply_elastic_trans:
            if self.prob_deform == 1:
                inputs = [
                    nrn_layers.SpatialTransformer(m)([v] + list_trans)
                    for (m, v) in zip(self.inter_method, inputs)
                ]
            else:
                rand_trans = tf.squeeze(
                    K.less(tf.random.uniform([1], 0, 1), self.prob_deform)
                )
                inputs = [
                    K.switch(
                        rand_trans,
                        nrn_layers.SpatialTransformer(m)([v] + list_trans),
                        v,
                    )
                    for (m, v) in zip(self.inter_method, inputs)
                ]
        if self.n_inputs < 2:
            return tf.cast(inputs[0], types[0])
        else:
            return [tf.cast(v, t) for (t, v) in zip(types, inputs)]


class RandomCrop(Layer):
    """Randomly crop all input tensors to a given shape. This cropping is applied to all channels.
    The input tensors are expected to have shape [batchsize, shape_dim1, ..., shape_dimn, channel].
    :param crop_shape: list with cropping shape in each dimension (excluding batch and channel dimension)

    example:
    if input is a tensor of shape [batchsize, 160, 160, 160, 3],
    output = RandomCrop(crop_shape=[96, 128, 96])(input)
    will yield an output of shape [batchsize, 96, 128, 96, 3] that is obtained by cropping with randomly selected
    cropping indices.
    """

    def __init__(self, crop_shape, **kwargs):

        self.several_inputs = True
        self.crop_max_val = None
        self.crop_shape = crop_shape
        self.n_dims = len(crop_shape)
        self.list_n_channels = None
        super(RandomCrop, self).__init__(**kwargs)

    def get_config(self):
        config = super().get_config()
        config["crop_shape"] = self.crop_shape
        return config

    def build(self, input_shape):

        if not isinstance(input_shape, list):
            self.several_inputs = False
            inputshape = [input_shape]
        else:
            inputshape = input_shape
        self.crop_max_val = np.array(
            np.array(inputshape[0][1 : self.n_dims + 1])
        ) - np.array(self.crop_shape)
        self.list_n_channels = [i[-1] for i in inputshape]
        self.built = True
        super(RandomCrop, self).build(input_shape)

    def call(self, inputs, **kwargs):

        # if one input only is provided, performs the cropping directly
        if not self.several_inputs:
            return tf.map_fn(self._single_slice, inputs, dtype=inputs.dtype)

        # otherwise we concatenate all inputs before cropping, so that they are all cropped at the same location
        else:
            types = [v.dtype for v in inputs]
            inputs = tf.concat([tf.cast(v, "float32") for v in inputs], axis=-1)
            inputs = tf.map_fn(self._single_slice, inputs, dtype=tf.float32)
            inputs = tf.split(inputs, self.list_n_channels, axis=-1)
            return [tf.cast(v, t) for (t, v) in zip(types, inputs)]

    def _single_slice(self, vol):
        crop_idx = tf.cast(
            tf.random.uniform([self.n_dims], 0, np.array(self.crop_max_val), "float32"),
            dtype="int32",
        )
        crop_idx = tf.concat([crop_idx, tf.zeros([1], dtype="int32")], axis=0)
        crop_size = tf.convert_to_tensor(self.crop_shape + [-1], dtype="int32")
        return tf.slice(vol, begin=crop_idx, size=crop_size)

    def compute_output_shape(self, input_shape):
        output_shape = [
            tuple([None] + self.crop_shape + [v]) for v in self.list_n_channels
        ]
        return output_shape if self.several_inputs else output_shape[0]


class RandomFlip(Layer):
    """This layer randomly flips the input tensor along the specified axes with a specified probability.
    It can also take multiple tensors as inputs (if they have the same shape). The same flips will be applied to all
    input tensors. These are expected to have shape [batchsize, shape_dim1, ..., shape_dimn, channel].
    If specified, this layer can also swap corresponding values. This is especially useful when flipping label maps
    with different labels for right/left structures, such that the flipped label maps keep a consistent labelling.
    :param axis: integer, or list of integers specifying the dimensions along which to flip.
    If a list, the input tensors can be flipped simultaneously in several directions. The values in flip_axis exclude
    the batch dimension (e.g. 0 will flip the tensor along the first axis after the batch dimension).
    Default is None, where the tensors can be flipped along all axes (except batch and channel axes).
    :param swap_labels: boolean to specify whether to swap the values of each input. Values are only swapped if an odd
    number of flips is applied.
    Can also be a list if several tensors are given as input.
    All the inputs for which the values need to be swapped must be int32 or int64.
    :param label_list: if swap_labels is True, list of all labels contained in labels. Must be ordered as follows, first
     the neutral labels (i.e. non-sided), then left labels and right labels.
    :param n_neutral_labels: if swap_labels is True, number of non-sided labels
    :param prob: probability to flip along each specified axis

    example 1:
    if input is a tensor of shape (batchsize, 10, 100, 200, 3)
    output = RandomFlip()(input) will randomly flip input along one of the 1st, 2nd, or 3rd axis (i.e. those with shape
    10, 100, 200).

    example 2:
    if input is a tensor of shape (batchsize, 10, 100, 200, 3)
    output = RandomFlip(flip_axis=1)(input) will randomly flip input along the 3rd axis (with shape 100), i.e. the axis
    with index 1 if we don't count the batch axis.

    example 3:
    input = tf.convert_to_tensor(np.array([[1, 0, 0, 0, 0, 0, 0],
                                           [1, 0, 0, 0, 2, 2, 0],
                                           [1, 0, 0, 0, 2, 2, 0],
                                           [1, 0, 0, 0, 2, 2, 0],
                                           [1, 0, 0, 0, 0, 0, 0]]))
    label_list = np.array([0, 1, 2])
    n_neutral_labels = 1
    output = RandomFlip(flip_axis=1, swap_labels=True, label_list=label_list, n_neutral_labels=n_neutral_labels)(input)
    where output will either be equal to input (bear in mind the flipping occurs with a 0.5 probability), or:
    output = [[0, 0, 0, 0, 0, 0, 2],
              [0, 1, 1, 0, 0, 0, 2],
              [0, 1, 1, 0, 0, 0, 2],
              [0, 1, 1, 0, 0, 0, 2],
              [0, 0, 0, 0, 0, 0, 2]]
    Note that the input must have a dtype int32 or int64 for its values to be swapped, otherwise an error will be raised

    example 4:
    if labels is the same as in the input of example 3, and image is a float32 image, then we can swap consistently both
    the labels and the image with:
    labels, image = RandomFlip(flip_axis=1, swap_labels=[True, False], label_list=label_list,
                               n_neutral_labels=n_neutral_labels)([labels, image]])
    Note that the labels must have a dtype int32 or int64 to be swapped, otherwise an error will be raised.
    This doesn't concern the image input, as its values are not swapped.
    """

    def __init__(
        self,
        axis=None,
        swap_labels=False,
        label_list=None,
        n_neutral_labels=None,
        prob=0.5,
        **kwargs,
    ):

        # shape attributes
        self.several_inputs = True
        self.n_dims = None
        self.list_n_channels = None

        # axis along which to flip
        self.axis = utils.reformat_to_list(axis)
        self.flip_axes = None

        # whether to swap labels, and corresponding label list
        self.swap_labels = utils.reformat_to_list(swap_labels)
        self.label_list = label_list
        self.n_neutral_labels = n_neutral_labels
        self.swap_lut = None

        self.prob = prob

        super(RandomFlip, self).__init__(**kwargs)

    def get_config(self):
        config = super().get_config()
        config["axis"] = self.axis
        config["swap_labels"] = self.swap_labels
        config["label_list"] = self.label_list
        config["n_neutral_labels"] = self.n_neutral_labels
        config["prob"] = self.prob
        return config

    def build(self, input_shape):

        if not isinstance(input_shape, list):
            self.several_inputs = False
            inputshape = [input_shape]
        else:
            inputshape = input_shape
        self.n_dims = len(inputshape[0][1:-1])
        self.list_n_channels = [i[-1] for i in inputshape]
        self.swap_labels = utils.reformat_to_list(
            self.swap_labels, length=len(inputshape)
        )
        self.flip_axes = (
            np.arange(self.n_dims).tolist() if self.axis is None else self.axis
        )

        # create label list with swapped labels
        if any(self.swap_labels):
            assert (self.label_list is not None) & (
                self.n_neutral_labels is not None
            ), "please provide a label_list, and n_neutral_labels when swapping the values of at least one input"
            n_labels = len(self.label_list)
            if self.n_neutral_labels == n_labels:
                self.swap_labels = [False] * len(self.swap_labels)
            else:
                rl_split = np.split(
                    self.label_list,
                    [
                        self.n_neutral_labels,
                        self.n_neutral_labels
                        + int((n_labels - self.n_neutral_labels) / 2),
                    ],
                )
                label_list_swap = np.concatenate(
                    (rl_split[0], rl_split[2], rl_split[1])
                )
                swap_lut = utils.get_mapping_lut(self.label_list, label_list_swap)
                self.swap_lut = tf.convert_to_tensor(swap_lut, dtype="int32")

        self.built = True
        super(RandomFlip, self).build(input_shape)

    def call(self, inputs, **kwargs):

        # convert inputs to list, and get each input type
        inputs = [inputs] if not self.several_inputs else inputs
        types = [v.dtype for v in inputs]

        # store whether to flip along each specified dimension
        batchsize = tf.split(tf.shape(inputs[0]), [1, self.n_dims + 1])[0]
        size = tf.concat(
            [batchsize, len(self.flip_axes) * tf.ones(1, dtype="int32")], axis=0
        )
        rand_flip = K.less(tf.random.uniform(size, 0, 1), self.prob)

        # swap right/left labels if we apply an odd number of flips
        odd = (
            tf.math.floormod(
                tf.reduce_sum(tf.cast(rand_flip, "int32"), -1, keepdims=True), 2
            )
            != 0
        )
        swapped_inputs = list()
        for i in range(len(inputs)):
            if self.swap_labels[i]:
                swapped_inputs.append(
                    tf.map_fn(self._single_swap, [inputs[i], odd], dtype=types[i])
                )
            else:
                swapped_inputs.append(inputs[i])

        # flip inputs and convert them back to their original type
        inputs = tf.concat([tf.cast(v, "float32") for v in swapped_inputs], axis=-1)
        inputs = tf.map_fn(self._single_flip, [inputs, rand_flip], dtype=tf.float32)
        inputs = tf.split(inputs, self.list_n_channels, axis=-1)

        if self.several_inputs:
            return [tf.cast(v, t) for (t, v) in zip(types, inputs)]
        else:
            return tf.cast(inputs[0], types[0])

    def _single_swap(self, inputs):
        return K.switch(inputs[1], tf.gather(self.swap_lut, inputs[0]), inputs[0])

    @staticmethod
    def _single_flip(inputs):
        flip_axis = tf.where(inputs[1])
        return K.switch(
            tf.equal(tf.size(flip_axis), 0),
            inputs[0],
            tf.reverse(inputs[0], axis=flip_axis[..., 0]),
        )


class SampleConditionalGMM(Layer):
    """This layer generates an image by sampling a Gaussian Mixture Model conditioned on a label map given as input.
    The parameters of the GMM are given as two additional inputs to the layer (means and standard deviations):
    image = SampleConditionalGMM(generation_labels)([label_map, means, stds])

    :param generation_labels: list of all possible label values contained in the input label maps.
    Must be a list or a 1D numpy array of size N, where N is the total number of possible label values.

    Layer inputs:
    label_map: input label map of shape [batchsize, shape_dim1, ..., shape_dimn, n_channel].
    All the values of label_map must be contained in generation_labels, but the input label_map doesn't necessarily have
    to contain all the values in generation_labels.
    means: tensor containing the mean values of all Gaussian distributions of the GMM.
           It must be of shape [batchsize, N, n_channel], and in the same order as generation label,
           i.e. the ith value of generation_labels will be associated to the ith value of means.
    stds: same as means but for the standard deviations of the GMM.
    """

    def __init__(self, generation_labels, **kwargs):
        self.generation_labels = generation_labels
        self.n_labels = None
        self.n_channels = None
        self.max_label = None
        self.indices = None
        self.shape = None
        super(SampleConditionalGMM, self).__init__(**kwargs)

    def get_config(self):
        config = super().get_config()
        config["generation_labels"] = self.generation_labels
        return config

    def build(self, input_shape):

        # check n_labels and n_channels
        assert (
            len(input_shape) == 3
        ), "should have three inputs: labels, means, std devs (in that order)."
        self.n_channels = input_shape[1][-1]
        self.n_labels = len(self.generation_labels)
        assert (
            self.n_labels == input_shape[1][1]
        ), "means should have the same number of values as generation_labels"
        assert (
            self.n_labels == input_shape[2][1]
        ), "stds should have the same number of values as generation_labels"

        # scatter parameters (to build mean/std lut)
        self.max_label = np.max(self.generation_labels) + 1
        indices = np.concatenate(
            [
                self.generation_labels + self.max_label * i
                for i in range(self.n_channels)
            ],
            axis=-1,
        )
        self.shape = tf.convert_to_tensor([np.max(indices) + 1], dtype="int32")
        self.indices = tf.convert_to_tensor(
            utils.add_axis(indices, axis=[0, -1]), dtype="int32"
        )

        self.built = True
        super(SampleConditionalGMM, self).build(input_shape)

    def call(self, inputs, **kwargs):

        # reformat labels and scatter indices
        batch = tf.split(tf.shape(inputs[0]), [1, -1])[0]
        tmp_indices = tf.tile(
            self.indices,
            tf.concat([batch, tf.convert_to_tensor([1, 1], dtype="int32")], axis=0),
        )
        labels = tf.concat(
            [
                tf.cast(inputs[0], dtype="int32") + self.max_label * i
                for i in range(self.n_channels)
            ],
            -1,
        )

        # build mean map
        means = tf.concat([inputs[1][..., i] for i in range(self.n_channels)], 1)
        tile_shape = tf.concat(
            [
                batch,
                tf.convert_to_tensor(
                    [
                        1,
                    ],
                    dtype="int32",
                ),
            ],
            axis=0,
        )
        means = tf.tile(
            tf.expand_dims(tf.scatter_nd(tmp_indices, means, self.shape), 0), tile_shape
        )
        means_map = tf.map_fn(
            lambda x: tf.gather(x[0], x[1]), [means, labels], dtype=tf.float32
        )

        # same for stds
        stds = tf.concat([inputs[2][..., i] for i in range(self.n_channels)], 1)
        stds = tf.tile(
            tf.expand_dims(tf.scatter_nd(tmp_indices, stds, self.shape), 0), tile_shape
        )
        stds_map = tf.map_fn(
            lambda x: tf.gather(x[0], x[1]), [stds, labels], dtype=tf.float32
        )

        return stds_map * tf.random.normal(tf.shape(labels)) + means_map

    def compute_output_shape(self, input_shape):
        return (
            input_shape[0]
            if (self.n_channels == 1)
            else tuple(list(input_shape[0][:-1]) + [self.n_channels])
        )


class SampleResolution(Layer):
    """Build a random resolution tensor by sampling a uniform distribution of provided range.

    You can use this layer in the following ways:
        resolution = SampleConditionalGMM(min_resolution)() in this case resolution will be a tensor of shape (n_dims,),
        where n_dims is the length of the min_resolution parameter (provided as a list, see below).
        resolution = SampleConditionalGMM(min_resolution)(input), where input is a tensor for which the first dimension
        represents the batch_size. In this case resolution will be a tensor of shape (batchsize, n_dims,).

    :param min_resolution: list of length n_dims specifying the inferior bounds of the uniform distributions to
    sample from for each value.
    :param max_res_iso: If not None, all the values of resolution will be equal to the same value, which is randomly
    sampled at each minibatch in U(min_resolution, max_res_iso).
    :param max_res_aniso: If not None, we first randomly select a direction i in the range [0, n_dims-1], and we sample
    a value in the corresponding uniform distribution U(min_resolution[i], max_res_aniso[i]).
    The other values of resolution will be set to min_resolution.
    :param prob_iso: if both max_res_iso and max_res_aniso are specified, this allows to specify the probability of
    sampling an isotropic resolution (therefore using max_res_iso) with respect to anisotropic resolution
    (which would use max_res_aniso).
    :param prob_min: if not zero, this allows to return with the specified probability an output resolution equal
    to min_resolution.
    :param return_thickness: if set to True, this layer will also return a thickness value of the same shape as
    resolution, which will be sampled independently for each axis from the uniform distribution
    U(min_resolution, resolution).

    """

    def __init__(
        self,
        min_resolution,
        max_res_iso=None,
        max_res_aniso=None,
        prob_iso=0.1,
        prob_min=0.05,
        return_thickness=True,
        **kwargs,
    ):

        self.min_res = min_resolution
        self.max_res_iso_input = max_res_iso
        self.max_res_iso = None
        self.max_res_aniso_input = max_res_aniso
        self.max_res_aniso = None
        self.prob_iso = prob_iso
        self.prob_min = prob_min
        self.return_thickness = return_thickness
        self.n_dims = len(self.min_res)
        self.add_batchsize = False
        self.min_res_tens = None
        super(SampleResolution, self).__init__(**kwargs)

    def get_config(self):
        config = super().get_config()
        config["min_resolution"] = self.min_res
        config["max_res_iso"] = self.max_res_iso
        config["max_res_aniso"] = self.max_res_aniso
        config["prob_iso"] = self.prob_iso
        config["prob_min"] = self.prob_min
        config["return_thickness"] = self.return_thickness
        return config

    def build(self, input_shape):

        # check maximum resolutions
        assert (self.max_res_iso_input is not None) | (
            self.max_res_aniso_input is not None
        ), "at least one of maximum isotropic or anisotropic resolutions must be provided, received none"

        # reformat resolutions as numpy arrays
        self.min_res = np.array(self.min_res)
        if self.max_res_iso_input is not None:
            self.max_res_iso = np.array(self.max_res_iso_input)
            assert len(self.min_res) == len(self.max_res_iso), (
                "min and isotropic max resolution must have the same length, "
                "had {0} and {1}".format(self.min_res, self.max_res_iso)
            )
            if np.array_equal(self.min_res, self.max_res_iso):
                self.max_res_iso = None
        if self.max_res_aniso_input is not None:
            self.max_res_aniso = np.array(self.max_res_aniso_input)
            assert len(self.min_res) == len(self.max_res_aniso), (
                "min and anisotropic max resolution must have the same length, "
                "had {} and {}".format(self.min_res, self.max_res_aniso)
            )
            if np.array_equal(self.min_res, self.max_res_aniso):
                self.max_res_aniso = None

        # check prob iso
        if (
            (self.max_res_iso is not None)
            & (self.max_res_aniso is not None)
            & (self.prob_iso == 0)
        ):
            raise Exception(
                "prob iso is 0 while sampling either isotropic and anisotropic resolutions is enabled"
            )

        if input_shape:
            self.add_batchsize = True

        self.min_res_tens = tf.convert_to_tensor(self.min_res, dtype="float32")

        self.built = True
        super(SampleResolution, self).build(input_shape)

    def call(self, inputs, **kwargs):

        if not self.add_batchsize:
            shape = [self.n_dims]
            dim = tf.random.uniform(
                shape=(1, 1), minval=0, maxval=self.n_dims, dtype="int32"
            )
            mask = tf.tensor_scatter_nd_update(
                tf.zeros([self.n_dims], dtype="bool"),
                dim,
                tf.convert_to_tensor([True], dtype="bool"),
            )
        else:
            batch = tf.split(tf.shape(inputs), [1, -1])[0]
            tile_shape = tf.concat(
                [batch, tf.convert_to_tensor([1], dtype="int32")], axis=0
            )
            self.min_res_tens = tf.tile(
                tf.expand_dims(self.min_res_tens, 0), tile_shape
            )

            shape = tf.concat(
                [batch, tf.convert_to_tensor([self.n_dims], dtype="int32")], axis=0
            )
            indices = tf.stack(
                [
                    tf.range(0, batch[0]),
                    tf.random.uniform(batch, 0, self.n_dims, dtype="int32"),
                ],
                1,
            )
            mask = tf.tensor_scatter_nd_update(
                tf.zeros(shape, dtype="bool"), indices, tf.ones(batch, dtype="bool")
            )

        # return min resolution as tensor if min=max
        if (self.max_res_iso is None) & (self.max_res_aniso is None):
            new_resolution = self.min_res_tens

        # sample isotropic resolution only
        elif (self.max_res_iso is not None) & (self.max_res_aniso is None):
            new_resolution_iso = tf.random.uniform(
                shape, minval=self.min_res, maxval=self.max_res_iso
            )
            new_resolution = K.switch(
                tf.squeeze(K.less(tf.random.uniform([1], 0, 1), self.prob_min)),
                self.min_res_tens,
                new_resolution_iso,
            )

        # sample anisotropic resolution only
        elif (self.max_res_iso is None) & (self.max_res_aniso is not None):
            new_resolution_aniso = tf.random.uniform(
                shape, minval=self.min_res, maxval=self.max_res_aniso
            )
            new_resolution = K.switch(
                tf.squeeze(K.less(tf.random.uniform([1], 0, 1), self.prob_min)),
                self.min_res_tens,
                tf.where(mask, new_resolution_aniso, self.min_res_tens),
            )

        # sample either anisotropic or isotropic resolution
        else:
            new_resolution_iso = tf.random.uniform(
                shape, minval=self.min_res, maxval=self.max_res_iso
            )
            new_resolution_aniso = tf.random.uniform(
                shape, minval=self.min_res, maxval=self.max_res_aniso
            )
            new_resolution = K.switch(
                tf.squeeze(K.less(tf.random.uniform([1], 0, 1), self.prob_iso)),
                new_resolution_iso,
                tf.where(mask, new_resolution_aniso, self.min_res_tens),
            )
            new_resolution = K.switch(
                tf.squeeze(K.less(tf.random.uniform([1], 0, 1), self.prob_min)),
                self.min_res_tens,
                new_resolution,
            )

        if self.return_thickness:
            return [
                new_resolution,
                tf.random.uniform(
                    tf.shape(self.min_res_tens), self.min_res_tens, new_resolution
                ),
            ]
        else:
            return new_resolution

    def compute_output_shape(self, input_shape):
        if self.return_thickness:
            return (
                [(None, self.n_dims)] * 2 if self.add_batchsize else [self.n_dims] * 2
            )
        else:
            return (None, self.n_dims) if self.add_batchsize else self.n_dims


class GaussianBlur(Layer):
    """Applies gaussian blur to an input image.
    The input image is expected to have shape [batchsize, shape_dim1, ..., shape_dimn, channel].
    :param sigma: standard deviation of the blurring kernels to apply. Can be a number, a list of length n_dims, or
    a numpy array.
    :param random_blur_range: (optional) if not None, this introduces a randomness in the blurring kernels, where
    sigma is now multiplied by a coefficient dynamically sampled from a uniform distribution with bounds
    [1/random_blur_range, random_blur_range].
    :param use_mask: (optional) whether a mask of the input will be provided as an additional layer input. This is used
    to mask the blurred image, and to correct for edge blurring effects.

    example 1:
    output = GaussianBlur(sigma=0.5)(input) will isotropically blur the input with a gaussian kernel of std 0.5.

    example 2:
    if input is a tensor of shape [batchsize, 10, 100, 200, 2]
    output = GaussianBlur(sigma=[0.5, 1, 10])(input) will blur the input a different gaussian kernel in each dimension.

    example 3:
    output = GaussianBlur(sigma=0.5, random_blur_range=1.15)(input)
    will blur the input a different gaussian kernel in each dimension, as each dimension will be associated with
    a kernel, whose standard deviation will be uniformly sampled from [0.5/1.15; 0.5*1.15].

    example 4:
    output = GaussianBlur(sigma=0.5, use_mask=True)([input, mask])
    will 1) blur the input a different gaussian kernel in each dimension, 2) mask the blurred image with the provided
    mask, and 3) correct for edge blurring effects. If the provided mask is not of boolean type, it will be thresholded
    above positive values.
    """

    def __init__(self, sigma, random_blur_range=None, use_mask=False, **kwargs):
        self.sigma = utils.reformat_to_list(sigma)
        assert np.all(
            np.array(self.sigma) >= 0
        ), "sigma should be superior or equal to 0"
        self.use_mask = use_mask

        self.n_dims = None
        self.n_channels = None
        self.blur_range = random_blur_range
        self.stride = None
        self.separable = None
        self.kernels = None
        self.convnd = None
        super(GaussianBlur, self).__init__(**kwargs)

    def get_config(self):
        config = super().get_config()
        config["sigma"] = self.sigma
        config["random_blur_range"] = self.blur_range
        config["use_mask"] = self.use_mask
        return config

    def build(self, input_shape):

        # get shapes
        if self.use_mask:
            assert (
                len(input_shape) == 2
            ), "please provide a mask as second layer input when use_mask=True"
            self.n_dims = len(input_shape[0]) - 2
            self.n_channels = input_shape[0][-1]
        else:
            self.n_dims = len(input_shape) - 2
            self.n_channels = input_shape[-1]

        # prepare blurring kernel
        self.stride = [1] * (self.n_dims + 2)
        self.sigma = utils.reformat_to_list(self.sigma, length=self.n_dims)
        self.separable = np.linalg.norm(np.array(self.sigma)) > 5
        if self.blur_range is None:  # fixed kernels
            self.kernels = l2i_et.gaussian_kernel(self.sigma, separable=self.separable)
        else:
            self.kernels = None

        # prepare convolution
        self.convnd = getattr(tf.nn, "conv%dd" % self.n_dims)

        self.built = True
        super(GaussianBlur, self).build(input_shape)

    def call(self, inputs, **kwargs):

        if self.use_mask:
            image = inputs[0]
            mask = tf.cast(inputs[1], "bool")
        else:
            image = inputs
            mask = None

        # redefine the kernels at each new step when blur_range is activated
        if self.blur_range is not None:
            self.kernels = l2i_et.gaussian_kernel(
                self.sigma, blur_range=self.blur_range, separable=self.separable
            )

        if self.separable:
            for k in self.kernels:
                if k is not None:
                    image = tf.concat(
                        [
                            self.convnd(
                                tf.expand_dims(image[..., n], -1),
                                k,
                                self.stride,
                                "SAME",
                            )
                            for n in range(self.n_channels)
                        ],
                        -1,
                    )
                    if self.use_mask:
                        maskb = tf.cast(mask, "float32")
                        maskb = tf.concat(
                            [
                                self.convnd(
                                    tf.expand_dims(maskb[..., n], -1),
                                    k,
                                    self.stride,
                                    "SAME",
                                )
                                for n in range(self.n_channels)
                            ],
                            -1,
                        )
                        image = image / (maskb + K.epsilon())
                        image = tf.where(mask, image, tf.zeros_like(image))
        else:
            if any(self.sigma):
                image = tf.concat(
                    [
                        self.convnd(
                            tf.expand_dims(image[..., n], -1),
                            self.kernels,
                            self.stride,
                            "SAME",
                        )
                        for n in range(self.n_channels)
                    ],
                    -1,
                )
                if self.use_mask:
                    maskb = tf.cast(mask, "float32")
                    maskb = tf.concat(
                        [
                            self.convnd(
                                tf.expand_dims(maskb[..., n], -1),
                                self.kernels,
                                self.stride,
                                "SAME",
                            )
                            for n in range(self.n_channels)
                        ],
                        -1,
                    )
                    image = image / (maskb + K.epsilon())
                    image = tf.where(mask, image, tf.zeros_like(image))

        return image


class DynamicGaussianBlur(Layer):
    """Applies gaussian blur to an input image, where the standard deviation of the blurring kernel is provided as a
    layer input, which enables to perform dynamic blurring (i.e. the blurring kernel can vary at each minibatch).
    :param max_sigma: maximum value of the standard deviation that will be provided as input. This is used to compute
    the size of the blurring kernels. It must be provided as a list of length n_dims.
    :param random_blur_range: (optional) if not None, this introduces a randomness in the blurring kernels, where
    sigma is now multiplied by a coefficient dynamically sampled from a uniform distribution with bounds
    [1/random_blur_range, random_blur_range].

    example:
    blurred_image = DynamicGaussianBlur(max_sigma=[5.]*3, random_blurring_range=1.15)([image, sigma])
    will return a blurred version of image, where the standard deviation of each dimension (given as a tensor, and with
    values lower than 5 for each axis) is multiplied by a random coefficient uniformly sampled from [1/1.15; 1.15].
    """

    def __init__(self, max_sigma, random_blur_range=None, **kwargs):
        self.max_sigma = max_sigma
        self.n_dims = None
        self.n_channels = None
        self.convnd = None
        self.blur_range = random_blur_range
        self.separable = None
        super(DynamicGaussianBlur, self).__init__(**kwargs)

    def get_config(self):
        config = super().get_config()
        config["max_sigma"] = self.max_sigma
        config["random_blur_range"] = self.blur_range
        return config

    def build(self, input_shape):
        assert (
            len(input_shape) == 2
        ), "sigma should be provided as an input tensor for dynamic blurring"
        self.n_dims = len(input_shape[0]) - 2
        self.n_channels = input_shape[0][-1]
        self.convnd = getattr(tf.nn, "conv%dd" % self.n_dims)
        self.max_sigma = utils.reformat_to_list(self.max_sigma, length=self.n_dims)
        self.separable = np.linalg.norm(np.array(self.max_sigma)) > 5
        self.built = True
        super(DynamicGaussianBlur, self).build(input_shape)

    def call(self, inputs, **kwargs):
        image = inputs[0]
        sigma = inputs[-1]
        kernels = l2i_et.gaussian_kernel(
            sigma, self.max_sigma, self.blur_range, self.separable
        )
        if self.separable:
            for kernel in kernels:
                image = tf.map_fn(self._single_blur, [image, kernel], dtype=tf.float32)
        else:
            image = tf.map_fn(self._single_blur, [image, kernels], dtype=tf.float32)
        return image

    def _single_blur(self, inputs):
        if self.n_channels > 1:
            split_channels = tf.split(inputs[0], [1] * self.n_channels, axis=-1)
            blurred_channel = list()
            for channel in split_channels:
                blurred = self.convnd(
                    tf.expand_dims(channel, 0),
                    inputs[1],
                    [1] * (self.n_dims + 2),
                    padding="SAME",
                )
                blurred_channel.append(tf.squeeze(blurred, axis=0))
            output = tf.concat(blurred_channel, -1)
        else:
            output = self.convnd(
                tf.expand_dims(inputs[0], 0),
                inputs[1],
                [1] * (self.n_dims + 2),
                padding="SAME",
            )
            output = tf.squeeze(output, axis=0)
        return output


class MimicAcquisition(Layer):
    """
    Layer that takes an image as input, and simulates data that has been acquired at low resolution.
    The output is obtained by resampling the input twice:
     - first at a resolution given as an input (i.e. the "acquisition" resolution),
     - then at the output resolution (specified output shape).
    The input tensor is expected to have shape [batchsize, shape_dim1, ..., shape_dimn, channel].

    :param volume_res: resolution of the provided inputs. Must be a 1-D numpy array with n_dims elements.
    :param min_subsample_res: lower bound of the acquisition resolutions to mimic (i.e. the input resolution must have
    values higher than min-subsample_res).
    :param resample_shape: shape of the output tensor
    :param build_dist_map: whether to return distance maps as outputs. These indicate the distance between each voxel
    and the nearest non-interpolated voxel (during the second resampling).
    :param prob_noise: probability to apply noise injection

    example 1:
    im_res = [1., 1., 1.]
    low_res = [1., 1., 3.]
    res = tf.convert_to_tensor([1., 1., 4.5])
    image is a tensor of shape (None, 256, 256, 256, 3)
    resample_shape = [256, 256, 256]
    output = MimicAcquisition(im_res, low_res, resample_shape)([image, res])
    output will be a tensor of shape (None, 256, 256, 256, 3), obtained by downsampling image to [1., 1., 4.5].
    and re-upsampling it at initial resolution (because resample_shape is equal to the input shape). In this example all
    examples of the batch will be downsampled to the same resolution (because res has no batch dimension).
    Note that the provided res must have higher values than min_low_res.

    example 2:
    im_res = [1., 1., 1.]
    min_low_res = [1., 1., 1.]
    res is a tensor of shape (None, 3), obtained for example by using the SampleResolution layer (see above).
    image is a tensor of shape (None, 256, 256, 256, 1)
    resample_shape = [128, 128, 128]
    output = MimicAcquisition(im_res, low_res, resample_shape)([image, res])
    output will be a tensor of shape (None, 128, 128, 128, 1), obtained by downsampling each examples of the batch to
    the matching resolution in res, and resampling them all to half the initial resolution.
    Note that the provided res must have higher values than min_low_res.
    """

    def __init__(
        self,
        volume_res,
        min_subsample_res,
        resample_shape,
        build_dist_map=False,
        noise_std=0,
        prob_noise=0.95,
        **kwargs,
    ):

        # resolutions and dimensions
        self.volume_res = volume_res
        self.min_subsample_res = min_subsample_res
        self.n_dims = len(self.volume_res)
        self.n_channels = None
        self.add_batchsize = None

        # noise
        self.noise_std = noise_std
        self.prob_noise = prob_noise

        # input and output shapes
        self.inshape = None
        self.resample_shape = resample_shape

        # meshgrids for resampling
        self.down_grid = None
        self.up_grid = None

        # whether to return a map indicating the distance from the interpolated voxels, to acquired ones.
        self.build_dist_map = build_dist_map

        super(MimicAcquisition, self).__init__(**kwargs)

    def get_config(self):
        config = super().get_config()
        config["volume_res"] = self.volume_res
        config["min_subsample_res"] = self.min_subsample_res
        config["resample_shape"] = self.resample_shape
        config["build_dist_map"] = self.build_dist_map
        config["noise_std"] = self.noise_std
        config["prob_noise"] = self.prob_noise
        return config

    def build(self, input_shape):

        # set up input shape and acquisition shape
        self.inshape = input_shape[0][1:]
        self.n_channels = input_shape[0][-1]
        self.add_batchsize = False if (input_shape[1][0] is None) else True
        down_tensor_shape = np.int32(
            np.array(self.inshape[:-1]) * self.volume_res / self.min_subsample_res
        )

        # build interpolation meshgrids
        self.down_grid = tf.expand_dims(
            tf.stack(nrn_utils.volshape_to_ndgrid(down_tensor_shape), -1), axis=0
        )
        self.up_grid = tf.expand_dims(
            tf.stack(nrn_utils.volshape_to_ndgrid(self.resample_shape), -1), axis=0
        )

        self.built = True
        super(MimicAcquisition, self).build(input_shape)

    def call(self, inputs, **kwargs):

        # sort inputs
        assert (
            len(inputs) == 2
        ), "inputs must have two items, the tensor to resample, and the downsampling resolution"
        vol = inputs[0]
        subsample_res = tf.cast(inputs[1], dtype="float32")
        vol = K.reshape(vol, [-1, *self.inshape])  # necessary for multi_gpu models
        batchsize = tf.split(tf.shape(vol), [1, -1])[0]
        tile_shape = tf.concat([batchsize, tf.ones([1], dtype="int32")], 0)

        # get downsampling and upsampling factors
        if self.add_batchsize:
            subsample_res = tf.tile(tf.expand_dims(subsample_res, 0), tile_shape)
        down_shape = tf.cast(
            tf.convert_to_tensor(
                np.array(self.inshape[:-1]) * self.volume_res, dtype="float32"
            )
            / subsample_res,
            dtype="int32",
        )
        down_zoom_factor = tf.cast(
            down_shape / tf.convert_to_tensor(self.inshape[:-1]), dtype="float32"
        )
        up_zoom_factor = tf.cast(
            tf.convert_to_tensor(self.resample_shape, dtype="int32") / down_shape,
            dtype="float32",
        )

        # downsample
        down_loc = tf.tile(
            self.down_grid,
            tf.concat([batchsize, tf.ones([self.n_dims + 1], dtype="int32")], 0),
        )
        down_loc = tf.cast(down_loc, "float32") / l2i_et.expand_dims(
            down_zoom_factor, axis=[1] * self.n_dims
        )
        inshape_tens = tf.tile(
            tf.expand_dims(tf.convert_to_tensor(self.inshape[:-1]), 0), tile_shape
        )
        inshape_tens = l2i_et.expand_dims(inshape_tens, axis=[1] * self.n_dims)
        down_loc = K.clip(down_loc, 0.0, tf.cast(inshape_tens, "float32"))
        vol = tf.map_fn(self._single_down_interpn, [vol, down_loc], tf.float32)

        # add noise with predefined probability
        if self.noise_std > 0:
            sample_shape = tf.concat(
                [
                    batchsize,
                    tf.ones([self.n_dims], dtype="int32"),
                    self.n_channels * tf.ones([1], dtype="int32"),
                ],
                0,
            )
            noise = tf.random.normal(
                tf.shape(vol),
                stddev=tf.random.uniform(sample_shape, maxval=self.noise_std),
            )
            if self.prob_noise == 1:
                vol += noise
            else:
                vol = K.switch(
                    tf.squeeze(K.less(tf.random.uniform([1], 0, 1), self.prob_noise)),
                    vol + noise,
                    vol,
                )

        # upsample
        up_loc = tf.tile(
            self.up_grid,
            tf.concat([batchsize, tf.ones([self.n_dims + 1], dtype="int32")], axis=0),
        )
        up_loc = tf.cast(up_loc, "float32") / l2i_et.expand_dims(
            up_zoom_factor, axis=[1] * self.n_dims
        )
        vol = tf.map_fn(self._single_up_interpn, [vol, up_loc], tf.float32)

        # return upsampled volume
        if not self.build_dist_map:
            return vol

        # return upsampled volumes with distance maps
        else:

            # get grid points
            floor = tf.math.floor(up_loc)
            ceil = tf.math.ceil(up_loc)

            # get distances of every voxel to higher and lower grid points for every dimension
            f_dist = up_loc - floor
            c_dist = ceil - up_loc

            # keep minimum 1d distances, and compute 3d distance to nearest grid point
            dist = tf.math.minimum(f_dist, c_dist) * l2i_et.expand_dims(
                subsample_res, axis=[1] * self.n_dims
            )
            dist = tf.math.sqrt(
                tf.math.reduce_sum(tf.math.square(dist), axis=-1, keepdims=True)
            )

            return [vol, dist]

    @staticmethod
    def _single_down_interpn(inputs):
        return nrn_utils.interpn(inputs[0], inputs[1], interp_method="nearest")

    @staticmethod
    def _single_up_interpn(inputs):
        return nrn_utils.interpn(inputs[0], inputs[1], interp_method="linear")

    def compute_output_shape(self, input_shape):
        output_shape = tuple([None] + self.resample_shape + [input_shape[0][-1]])
        return [output_shape] * 2 if self.build_dist_map else output_shape


class BiasFieldCorruption(Layer):
    """This layer applies a smooth random bias field to the input by applying the following steps:
    1) we first sample a value for the standard deviation of a centred normal distribution
    2) a small-size SVF is sampled from this normal distribution
    3) the small SVF is then resized with trilinear interpolation to image size
    4) it is rescaled to positive values by taking the voxel-wise exponential
    5) it is multiplied to the input tensor.
    The input tensor is expected to have shape [batchsize, shape_dim1, ..., shape_dimn, channel].

    :param bias_field_std: maximum value of the standard deviation sampled in 1 (it will be sampled from the range
    [0, bias_field_std])
    :param bias_scale: ratio between the shape of the input tensor and the shape of the sampled SVF.
    :param same_bias_for_all_channels: whether to apply the same bias field to all the channels of the input tensor.
    :param prob: probability to apply this bias field corruption.
    """

    def __init__(
        self,
        bias_field_std=0.5,
        bias_scale=0.025,
        same_bias_for_all_channels=False,
        prob=0.95,
        **kwargs,
    ):

        # input shape
        self.several_inputs = False
        self.inshape = None
        self.n_dims = None
        self.n_channels = None

        # sampling shape
        self.std_shape = None
        self.small_bias_shape = None

        # bias field parameters
        self.bias_field_std = bias_field_std
        self.bias_scale = bias_scale
        self.same_bias_for_all_channels = same_bias_for_all_channels
        self.prob = prob

        super(BiasFieldCorruption, self).__init__(**kwargs)

    def get_config(self):
        config = super().get_config()
        config["bias_field_std"] = self.bias_field_std
        config["bias_scale"] = self.bias_scale
        config["same_bias_for_all_channels"] = self.same_bias_for_all_channels
        config["prob"] = self.prob
        return config

    def build(self, input_shape):

        # input shape
        if isinstance(input_shape, list):
            self.several_inputs = True
            self.inshape = input_shape
        else:
            self.inshape = [input_shape]
        self.n_dims = len(self.inshape[0]) - 2
        self.n_channels = self.inshape[0][-1]

        # sampling shapes
        self.std_shape = [1] * (self.n_dims + 1)
        self.small_bias_shape = utils.get_resample_shape(
            self.inshape[0][1 : self.n_dims + 1], self.bias_scale, 1
        )
        if not self.same_bias_for_all_channels:
            self.std_shape[-1] = self.n_channels
            self.small_bias_shape[-1] = self.n_channels

        self.built = True
        super(BiasFieldCorruption, self).build(input_shape)

    def call(self, inputs, **kwargs):

        if not self.several_inputs:
            inputs = [inputs]

        if self.bias_field_std > 0:

            # sampling shapes
            batchsize = tf.split(tf.shape(inputs[0]), [1, -1])[0]
            std_shape = tf.concat(
                [batchsize, tf.convert_to_tensor(self.std_shape, dtype="int32")], 0
            )
            bias_shape = tf.concat(
                [batchsize, tf.convert_to_tensor(self.small_bias_shape, dtype="int32")],
                axis=0,
            )

            # sample small bias field
            bias_field = tf.random.normal(
                bias_shape,
                stddev=tf.random.uniform(std_shape, maxval=self.bias_field_std),
            )

            # resize bias field and take exponential
            bias_field = nrn_layers.Resize(
                size=self.inshape[0][1 : self.n_dims + 1], interp_method="linear"
            )(bias_field)
            bias_field = tf.math.exp(bias_field)

            # apply bias field with predefined probability
            if self.prob == 1:
                return [tf.math.multiply(bias_field, v) for v in inputs]
            else:
                rand_trans = tf.squeeze(K.less(tf.random.uniform([1], 0, 1), self.prob))
                if self.several_inputs:
                    return [
                        K.switch(rand_trans, tf.math.multiply(bias_field, v), v)
                        for v in inputs
                    ]
                else:
                    return K.switch(
                        rand_trans, tf.math.multiply(bias_field, inputs[0]), inputs[0]
                    )

        else:
            return inputs


class IntensityAugmentation(Layer):
    """This layer enables to augment the intensities of the input tensor, as well as to apply min_max normalisation.
    The following steps are applied (all are optional):
    1) white noise corruption, with a randomly sampled std dev.
    2) clip the input between two values
    3) min-max normalisation
    4) gamma augmentation (i.e. voxel-wise exponentiation by a randomly sampled power)
    The input tensor is expected to have shape [batchsize, shape_dim1, ..., shape_dimn, channel].

    :param noise_std: maximum value of the standard deviation of the Gaussian white noise used in 1 (it will be sampled
    from the range [0, noise_std]). Set to 0 to skip this step.
    :param clip: clip the input tensor between the given values. Can either be: a number (in which case we clip between
    0 and the given value), or a list or a numpy array with two elements. Default is 0, where no clipping occurs.
    :param normalise: whether to apply min-max normalisation, to normalise between 0 and 1. Default is True.
    :param norm_perc: percentiles (between 0 and 1) of the sorted intensity values for robust normalisation. Can be:
    a number (in which case the robust minimum is the provided percentile of sorted values, and the maximum is the
    1 - norm_perc percentile), or a list/numpy array of 2 elements (percentiles for the minimum and maximum values).
    The minimum and maximum values are computed separately for each channel if separate_channels is True.
    Default is 0, where we simply take the minimum and maximum values.
    :param gamma_std: standard deviation of the normal distribution from which we sample gamma (in log domain).
    Default is 0, where no gamma augmentation occurs.
    :param contrast_inversion: whether to perform contrast inversion (i.e. 1 - x). If True, this is performed randomly
    for each element of the batch, as well as for each channel.
    :param separate_channels: whether to augment all channels separately. Default is True.
    :param prob_noise: probability to apply noise injection
    :param prob_gamma: probability to apply gamma augmentation
    """

    def __init__(
        self,
        noise_std=0,
        clip=0,
        normalise=True,
        norm_perc=0,
        gamma_std=0,
        contrast_inversion=False,
        separate_channels=True,
        prob_noise=0.95,
        prob_gamma=1,
        **kwargs,
    ):

        # shape attributes
        self.n_dims = None
        self.n_channels = None
        self.flatten_shape = None
        self.expand_minmax_dim = None
        self.one = None

        # inputs
        self.noise_std = noise_std
        self.clip = clip
        self.clip_values = None
        self.normalise = normalise
        self.norm_perc = norm_perc
        self.perc = None
        self.gamma_std = gamma_std
        self.separate_channels = separate_channels
        self.contrast_inversion = contrast_inversion
        self.prob_noise = prob_noise
        self.prob_gamma = prob_gamma

        super(IntensityAugmentation, self).__init__(**kwargs)

    def get_config(self):
        config = super().get_config()
        config["noise_std"] = self.noise_std
        config["clip"] = self.clip
        config["normalise"] = self.normalise
        config["norm_perc"] = self.norm_perc
        config["gamma_std"] = self.gamma_std
        config["separate_channels"] = self.separate_channels
        config["prob_noise"] = self.prob_noise
        config["prob_gamma"] = self.prob_gamma
        return config

    def build(self, input_shape):
        self.n_dims = len(input_shape) - 2
        self.n_channels = input_shape[-1]
        self.flatten_shape = np.prod(np.array(input_shape[1:-1]))
        self.flatten_shape = (
            self.flatten_shape * self.n_channels
            if not self.separate_channels
            else self.flatten_shape
        )
        self.expand_minmax_dim = (
            self.n_dims if self.separate_channels else self.n_dims + 1
        )
        self.one = tf.ones([1], dtype="int32")
        if self.clip:
            self.clip_values = utils.reformat_to_list(self.clip)
            self.clip_values = (
                self.clip_values
                if len(self.clip_values) == 2
                else [0, self.clip_values[0]]
            )
        else:
            self.clip_values = None
        if self.norm_perc:
            self.perc = utils.reformat_to_list(self.norm_perc)
            self.perc = (
                self.perc if len(self.perc) == 2 else [self.perc[0], 1 - self.perc[0]]
            )
        else:
            self.perc = None

        self.built = True
        super(IntensityAugmentation, self).build(input_shape)

    def call(self, inputs, **kwargs):

        # prepare shape for sampling the noise and gamma std dev (depending on whether we augment channels separately)
        batchsize = tf.split(tf.shape(inputs), [1, -1])[0]
        if (self.noise_std > 0) | (self.gamma_std > 0) | self.contrast_inversion:
            sample_shape = tf.concat(
                [batchsize, tf.ones([self.n_dims], dtype="int32")], 0
            )
            if self.separate_channels:
                sample_shape = tf.concat([sample_shape, self.n_channels * self.one], 0)
            else:
                sample_shape = tf.concat([sample_shape, self.one], 0)
        else:
            sample_shape = None

        # add noise with predefined probability
        if self.noise_std > 0:
            noise_stddev = tf.random.uniform(sample_shape, maxval=self.noise_std)
            if self.separate_channels:
                noise = tf.random.normal(tf.shape(inputs), stddev=noise_stddev)
            else:
                noise = tf.random.normal(
                    tf.shape(tf.split(inputs, [1, -1], -1)[0]), stddev=noise_stddev
                )
                noise = tf.tile(
                    noise,
                    tf.convert_to_tensor([1] * (self.n_dims + 1) + [self.n_channels]),
                )
            if self.prob_noise == 1:
                inputs = inputs + noise
            else:
                inputs = K.switch(
                    tf.squeeze(K.less(tf.random.uniform([1], 0, 1), self.prob_noise)),
                    inputs + noise,
                    inputs,
                )

        # clip images to given values
        if self.clip_values is not None:
            inputs = K.clip(inputs, self.clip_values[0], self.clip_values[1])

        # normalise
        if self.normalise:
            # define robust min and max by sorting values and taking percentile
            if self.perc is not None:
                if self.separate_channels:
                    shape = tf.concat(
                        [
                            batchsize,
                            self.flatten_shape * self.one,
                            self.n_channels * self.one,
                        ],
                        0,
                    )
                else:
                    shape = tf.concat([batchsize, self.flatten_shape * self.one], 0)
                intensities = tf.sort(tf.reshape(inputs, shape), axis=1)
                m = intensities[:, max(int(self.perc[0] * self.flatten_shape), 0), ...]
                M = intensities[
                    :,
                    min(int(self.perc[1] * self.flatten_shape), self.flatten_shape - 1),
                    ...,
                ]
            # simple min and max
            else:
                m = K.min(inputs, axis=list(range(1, self.expand_minmax_dim + 1)))
                M = K.max(inputs, axis=list(range(1, self.expand_minmax_dim + 1)))
            # normalise
            m = l2i_et.expand_dims(m, axis=[1] * self.expand_minmax_dim)
            M = l2i_et.expand_dims(M, axis=[1] * self.expand_minmax_dim)
            inputs = tf.clip_by_value(inputs, m, M)
            inputs = (inputs - m) / (M - m + K.epsilon())

        # apply voxel-wise exponentiation with predefined probability
        if self.gamma_std > 0:
            gamma = tf.random.normal(sample_shape, stddev=self.gamma_std)
            if self.prob_gamma == 1:
                inputs = tf.math.pow(inputs, tf.math.exp(gamma))
            else:
                inputs = K.switch(
                    tf.squeeze(K.less(tf.random.uniform([1], 0, 1), self.prob_gamma)),
                    tf.math.pow(inputs, tf.math.exp(gamma)),
                    inputs,
                )

        # apply random contrast inversion
        if self.contrast_inversion:
            rand_invert = tf.less(tf.random.uniform(sample_shape, maxval=1), 0.5)
            split_channels = tf.split(inputs, [1] * self.n_channels, axis=-1)
            split_rand_invert = tf.split(rand_invert, [1] * self.n_channels, axis=-1)
            inverted_channel = list()
            for channel, invert in zip(split_channels, split_rand_invert):
                inverted_channel.append(
                    tf.map_fn(
                        self._single_invert, [channel, invert], dtype=channel.dtype
                    )
                )
            inputs = tf.concat(inverted_channel, -1)

        return inputs

    @staticmethod
    def _single_invert(inputs):
        return K.switch(tf.squeeze(inputs[1]), 1 - inputs[0], inputs[0])


class DiceLoss(Layer):
    """This layer computes the soft Dice loss between two tensors.
    These tensors are expected to have the same shape (one-hot encoding) [batch, size_dim1, ..., size_dimN, n_labels].
    The first input tensor is the GT and the second is the prediction: dice_loss = DiceLoss()([gt, pred])

    :param class_weights: (optional) if given, the loss is obtained by a weighted average of the Dice across labels.
    Must be a sequence or 1d numpy array of length n_labels. Can also be -1, where the weights are dynamically set to
    the inverse of the volume of each label in the ground truth.
    :param boundary_weights: (optional) bonus weight that we apply to the voxels close to boundaries between structures
    when computing the loss. Default is 0 where no boundary weighting is applied.
    :param boundary_dist: (optional) if boundary_weight is not 0, the extra boundary weighting is applied to all voxels
    within this distance to a region boundary. Default is 3.
    :param skip_background: (optional) whether to skip boundary weighting for the background class, as this may be
    redundant when we have several labels. This is only used if boundary_weight is not 0.
    :param enable_checks: (optional) whether to make sure that the 2 input tensors are probabilistic (i.e. the label
    probabilities sum to 1 at each voxel location). Default is True.
    """

    def __init__(
        self,
        class_weights=None,
        boundary_weights=0,
        boundary_dist=3,
        skip_background=True,
        enable_checks=True,
        **kwargs,
    ):

        self.class_weights = class_weights
        self.dynamic_weighting = False
        self.class_weights_tens = None
        self.boundary_weights = boundary_weights
        self.boundary_dist = boundary_dist
        self.skip_background = skip_background
        self.enable_checks = enable_checks
        self.spatial_axes = None
        self.avg_pooling_layer = None
        super(DiceLoss, self).__init__(**kwargs)

    def get_config(self):
        config = super().get_config()
        config["class_weights"] = self.class_weights
        config["boundary_weights"] = self.boundary_weights
        config["boundary_dist"] = self.boundary_dist
        config["skip_background"] = self.skip_background
        config["enable_checks"] = self.enable_checks
        return config

    def build(self, input_shape):

        # get shape
        assert (
            len(input_shape) == 2
        ), "DiceLoss expects 2 inputs to compute the Dice loss."
        assert (
            input_shape[0] == input_shape[1]
        ), "the two inputs must have the same shape."
        inshape = input_shape[0][1:]
        n_dims = len(inshape[:-1])
        n_labels = inshape[-1]
        self.spatial_axes = list(range(1, n_dims + 1))
        self.avg_pooling_layer = getattr(keras.layers, "AvgPool%dD" % n_dims)
        self.skip_background = False if n_labels == 1 else self.skip_background

        # build tensor with class weights
        if self.class_weights is not None:
            if self.class_weights == -1:
                self.dynamic_weighting = True
            else:
                class_weights_tens = utils.reformat_to_list(
                    self.class_weights, n_labels
                )
                class_weights_tens = tf.convert_to_tensor(class_weights_tens, "float32")
                self.class_weights_tens = l2i_et.expand_dims(class_weights_tens, 0)

        self.built = True
        super(DiceLoss, self).build(input_shape)

    def call(self, inputs, **kwargs):

        # make sure tensors are probabilistic
        gt = inputs[0]
        pred = inputs[1]
        if (
            self.enable_checks
        ):  # disabling is useful to, e.g., use incomplete label maps
            gt = K.clip(
                gt
                / (
                    tf.math.reduce_sum(gt, axis=-1, keepdims=True)
                    + tf.keras.backend.epsilon()
                ),
                0,
                1,
            )
            pred = K.clip(
                pred
                / (
                    tf.math.reduce_sum(pred, axis=-1, keepdims=True)
                    + tf.keras.backend.epsilon()
                ),
                0,
                1,
            )

        # compute dice loss for each label
        top = 2 * gt * pred
        bottom = tf.math.square(gt) + tf.math.square(pred)

        # apply boundary weighting (ie voxels close to region boundaries will be counted several times to compute Dice)
        if self.boundary_weights:
            avg = self.avg_pooling_layer(
                pool_size=2 * self.boundary_dist + 1, strides=1, padding="same"
            )(gt)
            boundaries = tf.cast(avg > 0.0, "float32") * tf.cast(
                avg < (1 / len(self.spatial_axes) - 1e-4), "float32"
            )
            if self.skip_background:
                boundaries_channels = tf.unstack(boundaries, axis=-1)
                boundaries = tf.stack(
                    [tf.zeros_like(boundaries_channels[0])] + boundaries_channels[1:],
                    axis=-1,
                )
            boundary_weights_tensor = 1 + self.boundary_weights * boundaries
            top *= boundary_weights_tensor
            bottom *= boundary_weights_tensor
        else:
            boundary_weights_tensor = None

        # compute loss
        top = tf.math.reduce_sum(top, self.spatial_axes)
        bottom = tf.math.reduce_sum(bottom, self.spatial_axes)
        dice = (top + tf.keras.backend.epsilon()) / (
            bottom + tf.keras.backend.epsilon()
        )
        loss = 1 - dice

        # apply class weighting across labels. In this case loss will have shape (batch), otherwise (batch, n_labels).
        if (
            self.dynamic_weighting
        ):  # the weight of a class is the inverse of its volume in the gt
            if (
                boundary_weights_tensor is not None
            ):  # we account for the boundary weighting to compute volume
                self.class_weights_tens = 1 / tf.reduce_sum(
                    gt * boundary_weights_tensor, self.spatial_axes
                )
            else:
                self.class_weights_tens = 1 / tf.reduce_sum(gt, self.spatial_axes)
        if self.class_weights_tens is not None:
            self.class_weights_tens /= tf.reduce_sum(self.class_weights_tens, -1)
            loss = tf.reduce_sum(loss * self.class_weights_tens, -1)

        return tf.math.reduce_mean(loss)

    def compute_output_shape(self, input_shape):
        return [[]]


class WeightedL2Loss(Layer):
    """This layer computes a L2 loss weighted by a specified factor (target_value) between two tensors.
    This is designed to be used on the layer before the softmax.
    The tensors are expected to have the same shape [batchsize, size_dim1, ..., size_dimN, n_labels].
    The first input tensor is the GT and the second is the prediction: wl2_loss = WeightedL2Loss()([gt, pred])

    :param target_value: target value for the layer before softmax: target_value when gt = 1, -target_value when gt = 0.
    """

    def __init__(self, target_value=5, **kwargs):
        self.target_value = target_value
        self.n_labels = None
        super(WeightedL2Loss, self).__init__(**kwargs)

    def get_config(self):
        config = super().get_config()
        config["target_value"] = self.target_value
        return config

    def build(self, input_shape):
        assert (
            len(input_shape) == 2
        ), "DiceLoss expects 2 inputs to compute the Dice loss."
        assert (
            input_shape[0] == input_shape[1]
        ), "the two inputs must have the same shape."
        self.n_labels = input_shape[0][-1]
        self.built = True
        super(WeightedL2Loss, self).build(input_shape)

    def call(self, inputs, **kwargs):
        gt = inputs[0]
        pred = inputs[1]
        weights = tf.expand_dims(1 - gt[..., 0] + 1e-8, -1)
        return K.sum(weights * K.square(pred - self.target_value * (2 * gt - 1))) / (
            K.sum(weights) * self.n_labels
        )

    def compute_output_shape(self, input_shape):
        return [[]]


class CrossEntropyLoss(Layer):
    """This layer computes the cross-entropy loss between two tensors.
    These tensors are expected to have the same shape (one-hot encoding) [batch, size_dim1, ..., size_dimN, n_labels].
    The first input tensor is the GT and the second is the prediction: ce_loss = CrossEntropyLoss()([gt, pred])

    :param class_weights: (optional) if given, the loss is obtained by a weighted average of the Dice across labels.
    Must be a sequence or 1d numpy array of length n_labels. Can also be -1, where the weights are dynamically set to
    the inverse of the volume of each label in the ground truth.
    :param boundary_weights: (optional) bonus weight that we apply to the voxels close to boundaries between structures
    when computing the loss. Default is 0 where no boundary weighting is applied.
    :param boundary_dist: (optional) if boundary_weight is not 0, the extra boundary weighting is applied to all voxels
    within this distance to a region boundary. Default is 3.
    :param skip_background: (optional) whether to skip boundary weighting for the background class, as this may be
    redundant when we have several labels. This is only used if boundary_weight is not 0.
    :param enable_checks: (optional) whether to make sure that the 2 input tensors are probabilistic (i.e. the label
    probabilities sum to 1 at each voxel location). Default is True.
    """

    def __init__(
        self,
        class_weights=None,
        boundary_weights=0,
        boundary_dist=3,
        skip_background=True,
        enable_checks=True,
        **kwargs,
    ):

        self.class_weights = class_weights
        self.dynamic_weighting = False
        self.class_weights_tens = None
        self.boundary_weights = boundary_weights
        self.boundary_dist = boundary_dist
        self.skip_background = skip_background
        self.enable_checks = enable_checks
        self.spatial_axes = None
        self.avg_pooling_layer = None
        super(CrossEntropyLoss, self).__init__(**kwargs)

    def get_config(self):
        config = super().get_config()
        config["class_weights"] = self.class_weights
        config["boundary_weights"] = self.boundary_weights
        config["boundary_dist"] = self.boundary_dist
        config["skip_background"] = self.skip_background
        config["enable_checks"] = self.enable_checks
        return config

    def build(self, input_shape):

        # get shape
        assert (
            len(input_shape) == 2
        ), "CrossEntropy expects 2 inputs to compute the Dice loss."
        assert (
            input_shape[0] == input_shape[1]
        ), "the two inputs must have the same shape."
        inshape = input_shape[0][1:]
        n_dims = len(inshape[:-1])
        n_labels = inshape[-1]
        self.spatial_axes = list(range(1, n_dims + 1))
        self.avg_pooling_layer = getattr(keras.layers, "AvgPool%dD" % n_dims)
        self.skip_background = False if n_labels == 1 else self.skip_background

        # build tensor with class weights
        if self.class_weights is not None:
            if self.class_weights == -1:
                self.dynamic_weighting = True
            else:
                class_weights_tens = utils.reformat_to_list(
                    self.class_weights, n_labels
                )
                class_weights_tens = tf.convert_to_tensor(class_weights_tens, "float32")
                self.class_weights_tens = l2i_et.expand_dims(
                    class_weights_tens, [0] * (1 + n_dims)
                )

        self.built = True
        super(CrossEntropyLoss, self).build(input_shape)

    def call(self, inputs, **kwargs):

        # make sure tensors are probabilistic
        gt = inputs[0]
        pred = inputs[1]
        if (
            self.enable_checks
        ):  # disabling is useful to, e.g., use incomplete label maps
            gt = K.clip(
                gt
                / (
                    tf.math.reduce_sum(gt, axis=-1, keepdims=True)
                    + tf.keras.backend.epsilon()
                ),
                0,
                1,
            )
            pred = pred / (
                tf.math.reduce_sum(pred, axis=-1, keepdims=True)
                + tf.keras.backend.epsilon()
            )
            pred = K.clip(
                pred, tf.keras.backend.epsilon(), 1 - tf.keras.backend.epsilon()
            )  # to avoid log(0)

        # compare prediction/target, ce has the same shape has the input tensors
        ce = -gt * tf.math.log(pred)

        # apply boundary weighting (ie voxels close to region boundaries will be counted several times to compute Dice)
        if self.boundary_weights:
            avg = self.avg_pooling_layer(
                pool_size=2 * self.boundary_dist + 1, strides=1, padding="same"
            )(gt)
            boundaries = tf.cast(avg > 0.0, "float32") * tf.cast(
                avg < (1 / len(self.spatial_axes) - 1e-4), "float32"
            )
            if self.skip_background:
                boundaries_channels = tf.unstack(boundaries, axis=-1)
                boundaries = tf.stack(
                    [tf.zeros_like(boundaries_channels[0])] + boundaries_channels[1:],
                    axis=-1,
                )
            boundary_weights_tensor = 1 + self.boundary_weights * boundaries
            ce *= boundary_weights_tensor
        else:
            boundary_weights_tensor = None

        # apply class weighting across labels. By the end of this, ce still has the same shape has the input tensors.
        if (
            self.dynamic_weighting
        ):  # the weight of a class is the inverse of its volume in the gt
            if (
                boundary_weights_tensor is not None
            ):  # we account for the boundary weighting to compute volume
                self.class_weights_tens = 1 / tf.reduce_sum(
                    gt * boundary_weights_tensor, self.spatial_axes, True
                )
            else:
                self.class_weights_tens = 1 / tf.reduce_sum(gt, self.spatial_axes)
        if self.class_weights_tens is not None:
            self.class_weights_tens /= tf.reduce_sum(self.class_weights_tens, -1)
            ce = tf.reduce_sum(ce * self.class_weights_tens, -1)

        # sum along label axis, and take the mean along spatial dimensions
        ce = tf.math.reduce_mean(tf.math.reduce_sum(ce, axis=-1))

        return ce

    def compute_output_shape(self, input_shape):
        return [[]]


class MomentLoss(Layer):
    """This layer computes a moment loss between two tensors. Specifically, it computes the distance between the centres
    of gravity for all the channels of the two tensors, and then returns a value averaged across all channels.
    These tensors are expected to have the same shape [batch, size_dim1, ..., size_dimN, n_channels].
    The first input tensor is the GT and the second is the prediction: moment_loss = MomentLoss()([gt, pred])

    :param class_weights: (optional) if given, the loss is obtained by a weighted average of the Dice across labels.
    Must be a sequence or 1d numpy array of length n_labels. Can also be -1, where the weights are dynamically set to
    the inverse of the volume of each label in the ground truth.
    :param enable_checks: (optional) whether to make sure that the 2 input tensors are probabilistic (i.e. the label
    probabilities sum to 1 at each voxel location). Default is True.
    """

    def __init__(self, class_weights=None, enable_checks=False, **kwargs):
        self.class_weights = class_weights
        self.dynamic_weighting = False
        self.class_weights_tens = None
        self.enable_checks = enable_checks
        self.spatial_axes = None
        self.coordinates = None
        super(MomentLoss, self).__init__(**kwargs)

    def get_config(self):
        config = super().get_config()
        config["class_weights"] = self.class_weights
        config["enable_checks"] = self.enable_checks
        return config

    def build(self, input_shape):

        # get shape
        assert (
            len(input_shape) == 2
        ), "MomentLoss expects 2 inputs to compute the Dice loss."
        assert (
            input_shape[0] == input_shape[1]
        ), "the two inputs must have the same shape."
        inshape = input_shape[0][1:]
        n_dims = len(inshape[:-1])
        n_labels = inshape[-1]
        self.spatial_axes = list(range(1, n_dims + 1))

        # build coordinate meshgrid of size (1, dim1, dim2, ..., dimN, ndim, nchan)
        self.coordinates = tf.stack(nrn_utils.volshape_to_ndgrid(inshape[:-1]), -1)
        self.coordinates = tf.cast(
            l2i_et.expand_dims(tf.stack([self.coordinates] * n_labels, -1), 0),
            "float32",
        )

        # build tensor with class weights
        if self.class_weights is not None:
            if self.class_weights == -1:
                self.dynamic_weighting = True
            else:
                class_weights_tens = utils.reformat_to_list(
                    self.class_weights, n_labels
                )
                class_weights_tens = tf.convert_to_tensor(class_weights_tens, "float32")
                self.class_weights_tens = l2i_et.expand_dims(class_weights_tens, 0)

        self.built = True
        super(MomentLoss, self).build(input_shape)

    def call(self, inputs, **kwargs):

        # make sure tensors are probabilistic
        gt = inputs[0]  # (B, dim1, dim2, ..., dimN, nchan)
        pred = inputs[1]
        if (
            self.enable_checks
        ):  # disabling is useful to, e.g., use incomplete label maps
            gt = gt / (
                tf.math.reduce_sum(gt, axis=-1, keepdims=True)
                + tf.keras.backend.epsilon()
            )
            pred = pred / (
                tf.math.reduce_sum(pred, axis=-1, keepdims=True)
                + tf.keras.backend.epsilon()
            )

        # compute loss
        gt_mean_coordinates = self._mean_coordinates(gt)  # (B, ndim, nchan)
        pred_mean_coordinates = self._mean_coordinates(pred)
        loss = tf.math.sqrt(
            tf.reduce_sum(
                tf.square(pred_mean_coordinates - gt_mean_coordinates), axis=1
            )
        )  # (B, nchan)

        # apply class weighting across labels. In this case loss will have shape (batch), otherwise (batch, n_labels).
        if (
            self.dynamic_weighting
        ):  # the weight of a class is the inverse of its volume in the gt
            self.class_weights_tens = 1 / tf.reduce_sum(gt, self.spatial_axes)
        if self.class_weights_tens is not None:
            self.class_weights_tens /= tf.reduce_sum(self.class_weights_tens, -1)
            loss = tf.reduce_sum(loss * self.class_weights_tens, -1)

        return tf.math.reduce_mean(loss)

    def _mean_coordinates(self, tensor):
        tensor = l2i_et.expand_dims(
            tensor, axis=-2
        )  # (B, dim1, dim2, ..., dimN, 1, nchan)
        numerator = tf.reduce_sum(
            tensor * self.coordinates, axis=self.spatial_axes
        )  # (B, ndim, nchan)
        denominator = (
            tf.reduce_sum(tensor, axis=self.spatial_axes) + tf.keras.backend.epsilon()
        )
        return numerator / denominator

    def compute_output_shape(self, input_shape):
        return [[]]


class ResetValuesToZero(Layer):
    """This layer enables to reset given values to 0 within the input tensors.

    :param values: list of values to be reset to 0.

    example:
    input = tf.convert_to_tensor(np.array([[1, 0, 2, 2, 2, 2, 0],
                                           [1, 3, 3, 3, 3, 3, 3],
                                           [1, 0, 0, 0, 4, 4, 4]]))
    values = [1, 3]
    ResetValuesToZero(values)(input)
    >> [[0, 0, 2, 2, 2, 2, 0],
        [0, 0, 0, 0, 0, 0, 0],
        [0, 0, 0, 0, 4, 4, 4]]
    """

    def __init__(self, values, **kwargs):
        assert (
            values is not None
        ), "please provide correct list of values, received None"
        self.values = utils.reformat_to_list(values)
        self.values_tens = None
        self.n_values = len(values)
        super(ResetValuesToZero, self).__init__(**kwargs)

    def get_config(self):
        config = super().get_config()
        config["values"] = self.values
        return config

    def build(self, input_shape):
        self.values_tens = tf.convert_to_tensor(self.values)
        self.built = True
        super(ResetValuesToZero, self).build(input_shape)

    def call(self, inputs, **kwargs):
        values = tf.cast(self.values_tens, dtype=inputs.dtype)
        for i in range(self.n_values):
            inputs = tf.where(
                tf.equal(inputs, values[i]), tf.zeros_like(inputs), inputs
            )
        return inputs


class ConvertLabels(Layer):
    """Convert all labels in a tensor by the corresponding given set of values.
    labels_converted = ConvertLabels(source_values, dest_values)(labels).
    labels must be an int32 tensor, and labels_converted will also be int32.

    :param source_values: list of all the possible values in labels. Must be a list or a 1D numpy array.
    :param dest_values: list of all the target label values. Must be ordered the same as source values:
    labels[labels == source_values[i]] = dest_values[i].
    If None (default), dest_values is equal to [0, ..., N-1], where N is the total number of values in source_values,
    which enables to remap label maps to [0, ..., N-1].
    """

    def __init__(self, source_values, dest_values=None, **kwargs):
        self.source_values = source_values
        self.dest_values = dest_values
        self.lut = None
        super(ConvertLabels, self).__init__(**kwargs)

    def get_config(self):
        config = super().get_config()
        config["source_values"] = self.source_values
        config["dest_values"] = self.dest_values
        return config

    def build(self, input_shape):
        self.lut = tf.convert_to_tensor(
            utils.get_mapping_lut(self.source_values, dest=self.dest_values),
            dtype="int32",
        )
        self.built = True
        super(ConvertLabels, self).build(input_shape)

    def call(self, inputs, **kwargs):
        return tf.gather(self.lut, tf.cast(inputs, dtype="int32"))


class PadAroundCentre(Layer):
    """Pad the input tensor to the specified shape with the given value.
    The input tensor is expected to have shape [batchsize, shape_dim1, ..., shape_dimn, channel].
    :param pad_margin: margin to use for padding. The tensor will be padded by the provided margin on each side.
    Can either be a number (all axes padded with the same margin), or a  list/numpy array of length n_dims.
    example: if tensor is of shape [batch, x, y, z, n_channels] and margin=10, then the padded tensor will be of
    shape [batch, x+2*10, y+2*10, z+2*10, n_channels].
    :param pad_shape: shape to pad the tensor to. Can either be a number (all axes padded to the same shape), or a
    list/numpy array of length n_dims.
    :param value: value to pad the tensors with. Default is 0.
    """

    def __init__(self, pad_margin=None, pad_shape=None, value=0, **kwargs):
        self.pad_margin = pad_margin
        self.pad_shape = pad_shape
        self.value = value
        self.pad_margin_tens = None
        self.pad_shape_tens = None
        self.n_dims = None
        super(PadAroundCentre, self).__init__(**kwargs)

    def get_config(self):
        config = super().get_config()
        config["pad_margin"] = self.pad_margin
        config["pad_shape"] = self.pad_shape
        config["value"] = self.value
        return config

    def build(self, input_shape):
        # input shape
        self.n_dims = len(input_shape) - 2
        shape = list(input_shape)
        shape[0] = 0
        shape[-1] = 0

        if self.pad_margin is not None:
            assert (
                self.pad_shape is None
            ), "please do not provide a padding shape and margin at the same time."

            # reformat padding margins
            pad = np.transpose(
                np.array(
                    [[0] + utils.reformat_to_list(self.pad_margin, self.n_dims) + [0]]
                    * 2
                )
            )
            self.pad_margin_tens = tf.convert_to_tensor(pad, dtype="int32")

        elif self.pad_shape is not None:
            assert (
                self.pad_margin is None
            ), "please do not provide a padding shape and margin at the same time."

            # pad shape
            tensor_shape = tf.cast(tf.convert_to_tensor(shape), "int32")
            self.pad_shape_tens = np.array(
                [0] + utils.reformat_to_list(self.pad_shape, length=self.n_dims) + [0]
            )
            self.pad_shape_tens = tf.convert_to_tensor(
                self.pad_shape_tens, dtype="int32"
            )
            self.pad_shape_tens = tf.math.maximum(tensor_shape, self.pad_shape_tens)

            # padding margin
            min_margins = (self.pad_shape_tens - tensor_shape) / 2
            max_margins = self.pad_shape_tens - tensor_shape - min_margins
            self.pad_margin_tens = tf.stack([min_margins, max_margins], axis=-1)

        else:
            raise Exception(
                "please either provide a padding shape or a padding margin."
            )

        self.built = True
        super(PadAroundCentre, self).build(input_shape)

    def call(self, inputs, **kwargs):
        return tf.pad(
            inputs, self.pad_margin_tens, mode="CONSTANT", constant_values=self.value
        )


class MaskEdges(Layer):
    """Reset the edges of a tensor to zero (i.e. with bands of zeros along the specified axes).
    The width of the zero-band is randomly drawn from a uniform distribution, whose range is given in boundaries.

    :param axes: axes along which to reset edges to zero. Can be an int (single axis), or a sequence.
    :param boundaries: numpy array of shape (len(axes), 4). Each row contains the two bounds of the uniform
    distributions from which we draw the width of the zero-bands on each side.
    Those bounds must be expressed in relative side (i.e. between 0 and 1).
    :return: a tensor of the same shape as the input, with bands of zeros along the specified axes.

    example:
    tensor=tf.constant([[[[1., 1., 1., 1., 1., 1., 1., 1., 1., 1.],
                       [1., 1., 1., 1., 1., 1., 1., 1., 1., 1.],
                       [1., 1., 1., 1., 1., 1., 1., 1., 1., 1.],
                       [1., 1., 1., 1., 1., 1., 1., 1., 1., 1.],
                       [1., 1., 1., 1., 1., 1., 1., 1., 1., 1.],
                       [1., 1., 1., 1., 1., 1., 1., 1., 1., 1.],
                       [1., 1., 1., 1., 1., 1., 1., 1., 1., 1.],
                       [1., 1., 1., 1., 1., 1., 1., 1., 1., 1.],
                       [1., 1., 1., 1., 1., 1., 1., 1., 1., 1.],
                       [1., 1., 1., 1., 1., 1., 1., 1., 1., 1.]]]])  # shape = [1,10,10,1]
    axes=1
    boundaries = np.array([[0.2, 0.45, 0.85, 0.9]])

    In this case, we reset the edges along the 2nd dimension (i.e. the 1st dimension after the batch dimension),
    the 1st zero-band will expand from the 1st row to a number drawn from [0.2*tensor.shape[1], 0.45*tensor.shape[1]],
    and the 2nd zero-band will expand from a row drawn from [0.85*tensor.shape[1], 0.9*tensor.shape[1]], to the end of
    the tensor. A possible output could be:
    array([[[[0., 0., 0., 0., 0., 0., 0., 0., 0., 0.],
          [0., 0., 0., 0., 0., 0., 0., 0., 0., 0.],
          [0., 0., 0., 0., 0., 0., 0., 0., 0., 0.],
          [1., 1., 1., 1., 1., 1., 1., 1., 1., 1.],
          [1., 1., 1., 1., 1., 1., 1., 1., 1., 1.],
          [1., 1., 1., 1., 1., 1., 1., 1., 1., 1.],
          [1., 1., 1., 1., 1., 1., 1., 1., 1., 1.],
          [1., 1., 1., 1., 1., 1., 1., 1., 1., 1.],
          [1., 1., 1., 1., 1., 1., 1., 1., 1., 1.],
          [0., 0., 0., 0., 0., 0., 0., 0., 0., 0.]]]])  # shape = [1,10,10,1]
    """

    def __init__(self, axes, boundaries, prob_mask=1, **kwargs):
        self.axes = utils.reformat_to_list(axes, dtype="int")
        self.boundaries = utils.reformat_to_n_channels_array(
            boundaries, n_dims=4, n_channels=len(self.axes)
        )
        self.prob_mask = prob_mask
        self.inputshape = None
        super(MaskEdges, self).__init__(**kwargs)

    def get_config(self):
        config = super().get_config()
        config["axes"] = self.axes
        config["boundaries"] = self.boundaries
        config["prob_mask"] = self.prob_mask
        return config

    def build(self, input_shape):
        self.inputshape = input_shape
        self.built = True
        super(MaskEdges, self).build(input_shape)

    def call(self, inputs, **kwargs):

        # build mask
        mask = tf.ones_like(inputs)
        for i, axis in enumerate(self.axes):

            # select restricting indices
            axis_boundaries = self.boundaries[i, :]
            idx1 = tf.math.round(
                tf.random.uniform(
                    [1],
                    minval=axis_boundaries[0] * self.inputshape[axis],
                    maxval=axis_boundaries[1] * self.inputshape[axis],
                )
            )
            idx2 = tf.math.round(
                tf.random.uniform(
                    [1],
                    minval=axis_boundaries[2] * self.inputshape[axis],
                    maxval=axis_boundaries[3] * self.inputshape[axis] - 1,
                )
                - idx1
            )
            idx3 = self.inputshape[axis] - idx1 - idx2
            split_idx = tf.cast(tf.concat([idx1, idx2, idx3], axis=0), dtype="int32")

            # update mask
            split_list = tf.split(inputs, split_idx, axis=axis)
            tmp_mask = tf.concat(
                [
                    tf.zeros_like(split_list[0]),
                    tf.ones_like(split_list[1]),
                    tf.zeros_like(split_list[2]),
                ],
                axis=axis,
            )
            mask = mask * tmp_mask

        # mask second_channel
        tensor = K.switch(
            tf.squeeze(K.greater(tf.random.uniform([1], 0, 1), 1 - self.prob_mask)),
            inputs * mask,
            inputs,
        )

        return [tensor, mask]

    def compute_output_shape(self, input_shape):
        return [input_shape] * 2


class ImageGradients(Layer):

    def __init__(self, gradient_type="sobel", return_magnitude=False, **kwargs):

        self.gradient_type = gradient_type
        assert (self.gradient_type == "sobel") | (
            self.gradient_type == "1-step_diff"
        ), (
            "gradient_type should be either sobel or 1-step_diff, had %s"
            % self.gradient_type
        )

        # shape
        self.n_dims = 0
        self.shape = None
        self.n_channels = 0

        # convolution params if sobel diff
        self.stride = None
        self.kernels = None
        self.convnd = None

        self.return_magnitude = return_magnitude

        super(ImageGradients, self).__init__(**kwargs)

    def get_config(self):
        config = super().get_config()
        config["gradient_type"] = self.gradient_type
        config["return_magnitude"] = self.return_magnitude
        return config

    def build(self, input_shape):

        # get shapes
        self.n_dims = len(input_shape) - 2
        self.shape = input_shape[1:]
        self.n_channels = input_shape[-1]

        # prepare kernel if sobel gradients
        if self.gradient_type == "sobel":
            self.kernels = l2i_et.sobel_kernels(self.n_dims)
            self.stride = [1] * (self.n_dims + 2)
            self.convnd = getattr(tf.nn, "conv%dd" % self.n_dims)
        else:
            self.kernels = self.convnd = self.stride = None

        self.built = True
        super(ImageGradients, self).build(input_shape)

    def call(self, inputs, **kwargs):

        image = inputs
        batchsize = tf.split(tf.shape(inputs), [1, -1])[0]
        gradients = list()

        # sobel method
        if self.gradient_type == "sobel":
            # get sobel gradients in each direction
            for n in range(self.n_dims):
                gradient = image
                # apply 1D kernel in each direction (sobel kernels are separable), instead of applying a nD kernel
                for k in self.kernels[n]:
                    gradient = tf.concat(
                        [
                            self.convnd(
                                tf.expand_dims(gradient[..., n], -1),
                                k,
                                self.stride,
                                "SAME",
                            )
                            for n in range(self.n_channels)
                        ],
                        -1,
                    )
                gradients.append(gradient)

        # 1-step method, only supports 2 and 3D
        else:

            # get 1-step diff
            if self.n_dims == 2:
                gradients.append(image[:, 1:, :, :] - image[:, :-1, :, :])  # dx
                gradients.append(image[:, :, 1:, :] - image[:, :, :-1, :])  # dy

            elif self.n_dims == 3:
                gradients.append(image[:, 1:, :, :, :] - image[:, :-1, :, :, :])  # dx
                gradients.append(image[:, :, 1:, :, :] - image[:, :, :-1, :, :])  # dy
                gradients.append(image[:, :, :, 1:, :] - image[:, :, :, :-1, :])  # dz

            else:
                raise Exception(
                    "ImageGradients only support 2D or 3D tensors for 1-step diff, had: %dD"
                    % self.n_dims
                )

            # pad with zeros to return tensors of the same shape as input
            for i in range(self.n_dims):
                tmp_shape = list(self.shape)
                tmp_shape[i] = 1
                zeros = tf.zeros(
                    tf.concat(
                        [batchsize, tf.convert_to_tensor(tmp_shape, dtype="int32")], 0
                    ),
                    image.dtype,
                )
                gradients[i] = tf.concat([gradients[i], zeros], axis=i + 1)

        # compute total gradient magnitude if necessary, or concatenate different gradients along the channel axis
        if self.return_magnitude:
            gradients = tf.sqrt(
                tf.reduce_sum(tf.square(tf.stack(gradients, axis=-1)), axis=-1)
            )
        else:
            gradients = tf.concat(gradients, axis=-1)

        return gradients

    def compute_output_shape(self, input_shape):
        if not self.return_magnitude:
            input_shape = list(input_shape)
            input_shape[-1] = self.n_dims
        return tuple(input_shape)


class RandomDilationErosion(Layer):
    """
    GPU implementation of binary dilation or erosion. The operation can be chosen to be always a dilation, or always an
    erosion, or randomly choosing between them for each element of the batch.
    The chosen operation is applied to the input with a given probability. Moreover, it is also possible to randomise
    the factor of the operation for each element of the mini-batch.
    :param min_factor: minimum possible value for the dilation/erosion factor. Must be an integer.
    :param max_factor: minimum possible value for the dilation/erosion factor. Must be an integer.
    Set it to the same value as min_factor to always perform dilation/erosion with the same factor.
    :param prob: probability with which to apply the selected operation to the input.
    :param operation: which operation to apply. Can be 'dilation' or 'erosion' or 'random'.
    :param return_mask: if operation is erosion and the input of this layer is a label map, we have the
    choice to either return the eroded label map or the mask (return_mask=True)
    """

    def __init__(
        self,
        min_factor,
        max_factor,
        max_factor_dilate=None,
        prob=1,
        operation="random",
        return_mask=False,
        **kwargs,
    ):

        self.min_factor = min_factor
        self.max_factor = max_factor
        self.max_factor_dilate = (
            max_factor_dilate if max_factor_dilate is not None else self.max_factor
        )
        self.prob = prob
        self.operation = operation
        self.return_mask = return_mask
        self.n_dims = None
        self.inshape = None
        self.n_channels = None
        self.convnd = None
        super(RandomDilationErosion, self).__init__(**kwargs)

    def get_config(self):
        config = super().get_config()
        config["min_factor"] = self.min_factor
        config["max_factor"] = self.max_factor
        config["max_factor_dilate"] = self.max_factor_dilate
        config["prob"] = self.prob
        config["operation"] = self.operation
        config["return_mask"] = self.return_mask
        return config

    def build(self, input_shape):

        # input shape
        self.inshape = input_shape
        self.n_dims = len(self.inshape) - 2
        self.n_channels = self.inshape[-1]

        # prepare convolution
        self.convnd = getattr(tf.nn, "conv%dd" % self.n_dims)

        self.built = True
        super(RandomDilationErosion, self).build(input_shape)

    def call(self, inputs, **kwargs):

        # sample probability of applying operation. If random negative is erosion and positive is dilation
        batchsize = tf.split(tf.shape(inputs), [1, -1])[0]
        shape = tf.concat([batchsize, tf.convert_to_tensor([1], dtype="int32")], axis=0)
        if self.operation == "dilation":
            prob = tf.random.uniform(shape, 0, 1)
        elif self.operation == "erosion":
            prob = tf.random.uniform(shape, -1, 0)
        elif self.operation == "random":
            prob = tf.random.uniform(shape, -1, 1)
        else:
            raise ValueError(
                "operation should either be 'dilation' 'erosion' or 'random', had %s"
                % self.operation
            )

        # build kernel
        if self.min_factor == self.max_factor:
            dist_threshold = self.min_factor * tf.ones(shape, dtype="int32")
        else:
            if (self.max_factor == self.max_factor_dilate) | (
                self.operation != "random"
            ):
                dist_threshold = tf.random.uniform(
                    shape, minval=self.min_factor, maxval=self.max_factor, dtype="int32"
                )
            else:
                dist_threshold = tf.cast(
                    tf.map_fn(self._sample_factor, [prob], dtype=tf.float32),
                    dtype="int32",
                )
        kernel = l2i_et.unit_kernel(
            dist_threshold, self.n_dims, max_dist_threshold=self.max_factor
        )

        # convolve input mask with kernel according to given probability
        mask = tf.cast(tf.cast(inputs, dtype="bool"), dtype="float32")
        mask = tf.map_fn(self._single_blur, [mask, kernel, prob], dtype=tf.float32)
        mask = tf.cast(mask, "bool")

        if self.return_mask:
            return mask
        else:
            return inputs * tf.cast(mask, dtype=inputs.dtype)

    def _sample_factor(self, inputs):
        return tf.cast(
            K.switch(
                K.less(tf.squeeze(inputs[0]), 0),
                tf.random.uniform(
                    (1,), self.min_factor, self.max_factor, dtype="int32"
                ),
                tf.random.uniform(
                    (1,), self.min_factor, self.max_factor_dilate, dtype="int32"
                ),
            ),
            dtype="float32",
        )

    def _single_blur(self, inputs):
        # dilate...
        new_mask = K.switch(
            K.greater(tf.squeeze(inputs[2]), 1 - self.prob + 0.001),
            tf.cast(
                tf.greater(
                    tf.squeeze(
                        self.convnd(
                            tf.expand_dims(inputs[0], 0),
                            inputs[1],
                            [1] * (self.n_dims + 2),
                            padding="SAME",
                        ),
                        axis=0,
                    ),
                    0.01,
                ),
                dtype="float32",
            ),
            inputs[0],
        )
        # ...or erode
        new_mask = K.switch(
            K.less(tf.squeeze(inputs[2]), -(1 - self.prob + 0.001)),
            1
            - tf.cast(
                tf.greater(
                    tf.squeeze(
                        self.convnd(
                            tf.expand_dims(1 - new_mask, 0),
                            inputs[1],
                            [1] * (self.n_dims + 2),
                            padding="SAME",
                        ),
                        axis=0,
                    ),
                    0.01,
                ),
                dtype="float32",
            ),
            new_mask,
        )
        return new_mask

    def compute_output_shape(self, input_shape):
        return input_shape
