"""

This file contains all the utilities used in that project. They are classified in 5 categories:

1- loading/saving functions:

    -load_volume

    -save_volume

    -get_volume_info

    -get_list_labels

    -load_array_if_path

    -write_pickle

    -read_pickle

    -write_model_summary

2- reformatting functions

    -reformat_to_list

    -reformat_to_n_channels_array

3- path related functions

    -list_images_in_folder

    -list_files

    -list_subfolders

    -strip_extension

    -strip_suffix

    -mkdir

    -mkcmd

4- shape-related functions

    -get_dims

    -get_resample_shape

    -add_axis

    -get_padding_margin

5- build affine matrices/tensors

    -create_affine_transformation_matrix

    -sample_affine_transform

    -create_rotation_transform

    -create_shearing_transform

6- miscellaneous

    -infer

    -LoopInfo

    -get_mapping_lut

    -build_training_generator

    -find_closest_number_divisible_by_m

    -build_binary_structure

    -draw_value_from_distribution

    -build_exp

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.

"""

import os

import glob

import math

import time

import pickle

import numpy as np

import nibabel as nib

import tensorflow as tf

import keras.layers as KL

import keras.backend as K

from datetime import timedelta

from scipy.ndimage.morphology import distance_transform_edt

# ---------------------------------------------- loading/saving functions ----------------------------------------------

def load_volume(path_volume, im_only=True, squeeze=True, dtype=None, aff_ref=None):

    """

    Load volume file.

    :param path_volume: path of the volume to load. Can either be a nii, nii.gz, mgz, or npz format.

    If npz format, 1) the variable name is assumed to be 'vol_data',

    2) the volume is associated with an identity affine matrix and blank header.

    :param im_only: (optional) if False, the function also returns the affine matrix and header of the volume.

    :param squeeze: (optional) whether to squeeze the volume when loading.

    :param dtype: (optional) if not None, convert the loaded volume to this numpy dtype.

    :param aff_ref: (optional) If not None, the loaded volume is aligned to this affine matrix.

    The returned affine matrix is also given in this new space. Must be a numpy array of dimension 4x4.

    :return: the volume, with corresponding affine matrix and header if im_only is False.

    """

    assert path_volume.endswith(('.nii', '.nii.gz', '.mgz', '.npz')), 'Unknown data file: %s' % path_volume

    if path_volume.endswith(('.nii', '.nii.gz', '.mgz')):

        x = nib.load(path_volume)

        if squeeze:

            volume = np.squeeze(x.get_fdata())

        else:

            volume = x.get_fdata()

        aff = x.affine

        header = x.header

    else:  # npz

        volume = np.load(path_volume)['vol_data']

        if squeeze:

            volume = np.squeeze(volume)

        aff = np.eye(4)

        header = nib.Nifti1Header()

    if dtype is not None:

        if 'int' in dtype:

            volume = np.round(volume)

        volume = volume.astype(dtype=dtype)

    # align image to reference affine matrix

    if aff_ref is not None:

        from ext.lab2im import edit_volumes  # the import is done here to avoid import loops

        n_dims, _ = get_dims(list(volume.shape), max_channels=10)

        volume, aff = edit_volumes.align_volume_to_ref(volume, aff, aff_ref=aff_ref, return_aff=True, n_dims=n_dims)

    if im_only:

        return volume

    else:

        return volume, aff, header

def save_volume(volume, aff, header, path, res=None, dtype=None, n_dims=3):

    """

    Save a volume.

    :param volume: volume to save

    :param aff: affine matrix of the volume to save. If aff is None, the volume is saved with an identity affine matrix.

    aff can also be set to 'FS', in which case the volume is saved with the affine matrix of FreeSurfer outputs.

    :param header: header of the volume to save. If None, the volume is saved with a blank header.

    :param path: path where to save the volume.

    :param res: (optional) update the resolution in the header before saving the volume.

    :param dtype: (optional) numpy dtype for the saved volume.

    :param n_dims: (optional) number of dimensions, to avoid confusion in multi-channel case. Default is None, where

    n_dims is automatically inferred.

    """

    mkdir(os.path.dirname(path))

    if '.npz' in path:

        np.savez_compressed(path, vol_data=volume)

    else:

        if header is None:

            header = nib.Nifti1Header()

        if isinstance(aff, str):

            if aff == 'FS':

                aff = np.array([[-1, 0, 0, 0], [0, 0, 1, 0], [0, -1, 0, 0], [0, 0, 0, 1]])

        elif aff is None:

            aff = np.eye(4)

        if dtype is not None:

            if 'int' in dtype:

                volume = np.round(volume)

            volume = volume.astype(dtype=dtype)

            nifty = nib.Nifti1Image(volume, aff, header)

            nifty.set_data_dtype(dtype)

        else:

            nifty = nib.Nifti1Image(volume, aff, header)

        if res is not None:

            if n_dims is None:

                n_dims, _ = get_dims(volume.shape)

            res = reformat_to_list(res, length=n_dims, dtype=None)

            nifty.header.set_zooms(res)

        nib.save(nifty, path)

def get_volume_info(path_volume, return_volume=False, aff_ref=None, max_channels=10):

    """

    Gather information about a volume: shape, affine matrix, number of dimensions and channels, header, and resolution.

    :param path_volume: path of the volume to get information form.

    :param return_volume: (optional) whether to return the volume along with the information.

    :param aff_ref: (optional) If not None, the loaded volume is aligned to this affine matrix.

    All info relative to the volume is then given in this new space. Must be a numpy array of dimension 4x4.

    :param max_channels: maximum possible number of channels for the input volume.

    :return: volume (if return_volume is true), and corresponding info. If aff_ref is not None, the returned aff is

    the original one, i.e. the affine of the image before being aligned to aff_ref.

    """

    # read image

    im, aff, header = load_volume(path_volume, im_only=False)

    # understand if image is multichannel

    im_shape = list(im.shape)

    n_dims, n_channels = get_dims(im_shape, max_channels=max_channels)

    im_shape = im_shape[:n_dims]

    # get labels res

    if '.nii' in path_volume:

        data_res = np.array(header['pixdim'][1:n_dims + 1])

    elif '.mgz' in path_volume:

        data_res = np.array(header['delta'])  # mgz image

    else:

        data_res = np.array([1.0] * n_dims)

    # align to given affine matrix

    if aff_ref is not None:

        from ext.lab2im import edit_volumes  # the import is done here to avoid import loops

        ras_axes = edit_volumes.get_ras_axes(aff, n_dims=n_dims)

        ras_axes_ref = edit_volumes.get_ras_axes(aff_ref, n_dims=n_dims)

        im = edit_volumes.align_volume_to_ref(im, aff, aff_ref=aff_ref, n_dims=n_dims)

        im_shape = np.array(im_shape)

        data_res = np.array(data_res)

        im_shape[ras_axes_ref] = im_shape[ras_axes]

        data_res[ras_axes_ref] = data_res[ras_axes]

        im_shape = im_shape.tolist()

    # return info

    if return_volume:

        return im, im_shape, aff, n_dims, n_channels, header, data_res

    else:

        return im_shape, aff, n_dims, n_channels, header, data_res

def get_list_labels(label_list=None, labels_dir=None, save_label_list=None, FS_sort=False):

    """This function reads or computes a list of all label values used in a set of label maps.

    It can also sort all labels according to FreeSurfer lut.

    :param label_list: (optional) already computed label_list. Can be a sequence, a 1d numpy array, or the path to

    a numpy 1d array.

    :param labels_dir: (optional) if path_label_list is None, the label list is computed by reading all the label maps

    in the given folder. Can also be the path to a single label map.

    :param save_label_list: (optional) path where to save the label list.

    :param FS_sort: (optional) whether to sort label values according to the FreeSurfer classification.

    If true, the label values will be ordered as follows: neutral labels first (i.e. non-sided), left-side labels,

    and right-side labels. If FS_sort is True, this function also returns the number of neutral labels in label_list.

    :return: the label list (numpy 1d array), and the number of neutral (i.e. non-sided) labels if FS_sort is True.

    If one side of the brain is not represented at all in label_list, all labels are considered as neutral, and

    n_neutral_labels = len(label_list).

    """

    # load label list if previously computed

    if label_list is not None:

        label_list = np.array(reformat_to_list(label_list, load_as_numpy=True, dtype='int'))

    # compute label list from all label files

    elif labels_dir is not None:

        print('Compiling list of unique labels')

        # go through all labels files and compute unique list of labels

        labels_paths = list_images_in_folder(labels_dir)

        label_list = np.empty(0)

        loop_info = LoopInfo(len(labels_paths), 10, 'processing', print_time=True)

        for lab_idx, path in enumerate(labels_paths):

            loop_info.update(lab_idx)

            y = load_volume(path, dtype='int32')

            y_unique = np.unique(y)

            label_list = np.unique(np.concatenate((label_list, y_unique))).astype('int')

    else:

        raise Exception('either label_list, path_label_list or labels_dir should be provided')

    # sort labels in neutral/left/right according to FS labels

    n_neutral_labels = 0

    if FS_sort:

        neutral_FS_labels = [0, 14, 15, 16, 21, 22, 23, 24, 72, 77, 80, 85, 100, 101, 102, 103, 104, 105, 106, 107, 108,

                             109, 165, 200, 201, 202, 203, 204, 205, 206, 207, 208, 209, 210,

                             251, 252, 253, 254, 255, 258, 259, 260, 331, 332, 333, 334, 335, 336, 337, 338, 339, 340,

                             502, 506, 507, 508, 509, 511, 512, 514, 515, 516, 517, 530,

                             531, 532, 533, 534, 535, 536, 537]

        neutral = list()

        left = list()

        right = list()

        for la in label_list:

            if la in neutral_FS_labels:

                if la not in neutral:

                    neutral.append(la)

            elif (0 < la < 14) | (16 < la < 21) | (24 < la < 40) | (135 < la < 139) | (1000 <= la <= 1035) | \

                    (la == 865) | (20100 < la < 20110):

                if la not in left:

                    left.append(la)

            elif (39 < la < 72) | (162 < la < 165) | (2000 <= la <= 2035) | (20000 < la < 20010) | (la == 139) | \

                    (la == 866):

                if la not in right:

                    right.append(la)

            else:

                raise Exception('label {} not in our current FS classification, '

                                'please update get_list_labels in utils.py'.format(la))

        label_list = np.concatenate([sorted(neutral), sorted(left), sorted(right)])

        if ((len(left) > 0) & (len(right) > 0)) | ((len(left) == 0) & (len(right) == 0)):

            n_neutral_labels = len(neutral)

        else:

            n_neutral_labels = len(label_list)

    # save labels if specified

    if save_label_list is not None:

        np.save(save_label_list, np.int32(label_list))

    if FS_sort:

        return np.int32(label_list), n_neutral_labels

    else:

        return np.int32(label_list), None

def load_array_if_path(var, load_as_numpy=True):

    """If var is a string and load_as_numpy is True, this function loads the array writen at the path indicated by var.

    Otherwise it simply returns var as it is."""

    if (isinstance(var, str)) & load_as_numpy:

        assert os.path.isfile(var), 'No such path: %s' % var

        var = np.load(var)

    return var

def write_pickle(filepath, obj):

    """ write a python object with a pickle at a given path"""

    with open(filepath, 'wb') as file:

        pickler = pickle.Pickler(file)

        pickler.dump(obj)

def read_pickle(filepath):

    """ read a python object with a pickle"""

    with open(filepath, 'rb') as file:

        unpickler = pickle.Unpickler(file)

        return unpickler.load()

def write_model_summary(model, filepath='./model_summary.txt', line_length=150):

    """Write the summary of a keras model at a given path, with a given length for each line"""

    with open(filepath, 'w') as fh:

        model.summary(print_fn=lambda x: fh.write(x + '\n'), line_length=line_length)

# ----------------------------------------------- reformatting functions -----------------------------------------------

def reformat_to_list(var, length=None, load_as_numpy=False, dtype=None):

    """This function takes a variable and reformat it into a list of desired

    length and type (int, float, bool, str).

    If variable is a string, and load_as_numpy is True, it will be loaded as a numpy array.

    If variable is None, this function returns None.

    :param var: a str, int, float, list, tuple, or numpy array

    :param length: (optional) if var is a single item, it will be replicated to a list of this length

    :param load_as_numpy: (optional) whether var is the path to a numpy array

    :param dtype: (optional) convert all item to this type. Can be 'int', 'float', 'bool', or 'str'

    :return: reformatted list

    """

    # convert to list

    if var is None:

        return None

    var = load_array_if_path(var, load_as_numpy=load_as_numpy)

    if isinstance(var, (int, float, np.int, np.int32, np.int64, np.float, np.float32, np.float64)):

        var = [var]

    elif isinstance(var, tuple):

        var = list(var)

    elif isinstance(var, np.ndarray):

        if var.shape == (1,):

            var = [var[0]]

        else:

            var = np.squeeze(var).tolist()

    elif isinstance(var, str):

        var = [var]

    elif isinstance(var, bool):

        var = [var]

    if isinstance(var, list):

        if length is not None:

            if len(var) == 1:

                var = var * length

            elif len(var) != length:

                raise ValueError('if var is a list/tuple/numpy array, it should be of length 1 or {0}, '

                                 'had {1}'.format(length, var))

    else:

        raise TypeError('var should be an int, float, tuple, list, numpy array, or path to numpy array')

    # convert items type

    if dtype is not None:

        if dtype == 'int':

            var = [int(v) for v in var]

        elif dtype == 'float':

            var = [float(v) for v in var]

        elif dtype == 'bool':

            var = [bool(v) for v in var]

        elif dtype == 'str':

            var = [str(v) for v in var]

        else:

            raise ValueError("dtype should be 'str', 'float', 'int', or 'bool'; had {}".format(dtype))

    return var

def reformat_to_n_channels_array(var, n_dims=3, n_channels=1):

    """This function takes an int, float, list or tuple and reformat it to an array of shape (n_channels, n_dims).

    If resolution is a str, it will be assumed to be the path of a numpy array.

    If resolution is a numpy array, it will be checked to have shape (n_channels, n_dims).

    Finally if resolution is None, this function returns None as well."""

    if var is None:

        return [None] * n_channels

    if isinstance(var, str):

        var = np.load(var)

    # convert to numpy array

    if isinstance(var, (int, float, list, tuple)):

        var = reformat_to_list(var, n_dims)

        var = np.tile(np.array(var), (n_channels, 1))

    # check shape if numpy array

    elif isinstance(var, np.ndarray):

        if n_channels == 1:

            var = var.reshape((1, n_dims))

        else:

            if np.squeeze(var).shape == (n_dims,):

                var = np.tile(var.reshape((1, n_dims)), (n_channels, 1))

            elif var.shape != (n_channels, n_dims):

                raise ValueError('if array, var should be {0} or {1}'.format((1, n_dims), (n_channels, n_dims)))

    else:

        raise TypeError('var should be int, float, list, tuple or ndarray')

    return np.round(var, 3)

# ----------------------------------------------- path-related functions -----------------------------------------------

def list_images_in_folder(path_dir, include_single_image=True, check_if_empty=True):

    """List all files with extension nii, nii.gz, mgz, or npz within a folder."""

    basename = os.path.basename(path_dir)

    if include_single_image & \

            (('.nii.gz' in basename) | ('.nii' in basename) | ('.mgz' in basename) | ('.npz' in basename)):

        assert os.path.isfile(path_dir), 'file %s does not exist' % path_dir

        list_images = [path_dir]

    else:

        if os.path.isdir(path_dir):

            list_images = sorted(glob.glob(os.path.join(path_dir, '*nii.gz')) +

                                 glob.glob(os.path.join(path_dir, '*nii')) +

                                 glob.glob(os.path.join(path_dir, '*.mgz')) +

                                 glob.glob(os.path.join(path_dir, '*.npz')))

        else:

            raise Exception('Folder does not exist: %s' % path_dir)

        if check_if_empty:

            assert len(list_images) > 0, 'no .nii, .nii.gz, .mgz or .npz image could be found in %s' % path_dir

    return list_images

def list_files(path_dir, whole_path=True, expr=None, cond_type='or'):

    """This function returns a list of files contained in a folder, with possible regexp.

    :param path_dir: path of a folder

    :param whole_path: (optional) whether to return whole path or just the filenames.

    :param expr: (optional) regexp for files to list. Can be a str or a list of str.

    :param cond_type: (optional) if exp is a list, specify the logical link between expressions in exp.

    Can be 'or', or 'and'.

    :return: a list of files

    """

    assert isinstance(whole_path, bool), "whole_path should be bool"

    assert cond_type in ['or', 'and'], "cond_type should be either 'or', or 'and'"

    if whole_path:

        files_list = sorted([os.path.join(path_dir, f) for f in os.listdir(path_dir)

                             if os.path.isfile(os.path.join(path_dir, f))])

    else:

        files_list = sorted([f for f in os.listdir(path_dir) if os.path.isfile(os.path.join(path_dir, f))])

    if expr is not None:  # assumed to be either str or list of str

        if isinstance(expr, str):

            expr = [expr]

        elif not isinstance(expr, (list, tuple)):

            raise Exception("if specified, 'expr' should be a string or list of strings.")

        matched_list_files = list()

        for match in expr:

            tmp_matched_files_list = sorted([f for f in files_list if match in os.path.basename(f)])

            if cond_type == 'or':

                files_list = [f for f in files_list if f not in tmp_matched_files_list]

                matched_list_files += tmp_matched_files_list

            elif cond_type == 'and':

                files_list = tmp_matched_files_list

                matched_list_files = tmp_matched_files_list

        files_list = sorted(matched_list_files)

    return files_list

def list_subfolders(path_dir, whole_path=True, expr=None, cond_type='or'):

    """This function returns a list of subfolders contained in a folder, with possible regexp.

    :param path_dir: path of a folder

    :param whole_path: (optional) whether to return whole path or just the subfolder names.

    :param expr: (optional) regexp for files to list. Can be a str or a list of str.

    :param cond_type: (optional) if exp is a list, specify the logical link between expressions in exp.

    Can be 'or', or 'and'.

    :return: a list of subfolders

    """

    assert isinstance(whole_path, bool), "whole_path should be bool"

    assert cond_type in ['or', 'and'], "cond_type should be either 'or', or 'and'"

    if whole_path:

        subdirs_list = sorted([os.path.join(path_dir, f) for f in os.listdir(path_dir)

                               if os.path.isdir(os.path.join(path_dir, f))])

    else:

        subdirs_list = sorted([f for f in os.listdir(path_dir) if os.path.isdir(os.path.join(path_dir, f))])

    if expr is not None:  # assumed to be either str or list of str

        if isinstance(expr, str):

            expr = [expr]

        elif not isinstance(expr, (list, tuple)):

            raise Exception("if specified, 'expr' should be a string or list of strings.")

        matched_list_subdirs = list()

        for match in expr:

            tmp_matched_list_subdirs = sorted([f for f in subdirs_list if match in os.path.basename(f)])

            if cond_type == 'or':

                subdirs_list = [f for f in subdirs_list if f not in tmp_matched_list_subdirs]

                matched_list_subdirs += tmp_matched_list_subdirs

            elif cond_type == 'and':

                subdirs_list = tmp_matched_list_subdirs

                matched_list_subdirs = tmp_matched_list_subdirs

        subdirs_list = sorted(matched_list_subdirs)

    return subdirs_list

def get_image_extension(path):

    name = os.path.basename(path)

    if name[-7:] == '.nii.gz':

        return 'nii.gz'

    elif name[-4:] == '.mgz':

        return 'mgz'

    elif name[-4:] == '.nii':

        return 'nii'

    elif name[-4:] == '.npz':

        return 'npz'

def strip_extension(path):

    """Strip classical image extensions (.nii.gz, .nii, .mgz, .npz) from a filename."""

    return path.replace('.nii.gz', '').replace('.nii', '').replace('.mgz', '').replace('.npz', '')

def strip_suffix(path):

    """Strip classical image suffix from a filename."""

    path = path.replace('_aseg', '')

    path = path.replace('aseg', '')

    path = path.replace('.aseg', '')

    path = path.replace('_aseg_1', '')

    path = path.replace('_aseg_2', '')

    path = path.replace('aseg_1_', '')

    path = path.replace('aseg_2_', '')

    path = path.replace('_orig', '')

    path = path.replace('orig', '')

    path = path.replace('.orig', '')

    path = path.replace('_norm', '')

    path = path.replace('norm', '')

    path = path.replace('.norm', '')

    path = path.replace('_talairach', '')

    path = path.replace('GSP_FS_4p5', 'GSP')

    path = path.replace('.nii_crispSegmentation', '')

    path = path.replace('_crispSegmentation', '')

    path = path.replace('_seg', '')

    path = path.replace('.seg', '')

    path = path.replace('seg', '')

    path = path.replace('_seg_1', '')

    path = path.replace('_seg_2', '')

    path = path.replace('seg_1_', '')

    path = path.replace('seg_2_', '')

    return path

def mkdir(path_dir):

    """Recursively creates the current dir as well as its parent folders if they do not already exist."""

    if path_dir[-1] == '/':

        path_dir = path_dir[:-1]

    if not os.path.isdir(path_dir):

        list_dir_to_create = [path_dir]

        while not os.path.isdir(os.path.dirname(list_dir_to_create[-1])):

            list_dir_to_create.append(os.path.dirname(list_dir_to_create[-1]))

        for dir_to_create in reversed(list_dir_to_create):

            os.mkdir(dir_to_create)

def mkcmd(*args):

    """Creates terminal command with provided inputs.

    Example: mkcmd('mv', 'source', 'dest') will give 'mv source dest'."""

    return ' '.join([str(arg) for arg in args])

# ---------------------------------------------- shape-related functions -----------------------------------------------

def get_dims(shape, max_channels=10):

    """Get the number of dimensions and channels from the shape of an array.

    The number of dimensions is assumed to be the length of the shape, as long as the shape of the last dimension is

    inferior or equal to max_channels (default 3).

    :param shape: shape of an array. Can be a sequence or a 1d numpy array.

    :param max_channels: maximum possible number of channels.

    :return: the number of dimensions and channels associated with the provided shape.

    example 1: get_dims([150, 150, 150], max_channels=10) = (3, 1)

    example 2: get_dims([150, 150, 150, 3], max_channels=10) = (3, 3)

    example 3: get_dims([150, 150, 150, 15], max_channels=10) = (4, 1), because 5>3"""

    if shape[-1] <= max_channels:

        n_dims = len(shape) - 1

        n_channels = shape[-1]

    else:

        n_dims = len(shape)

        n_channels = 1

    return n_dims, n_channels

def get_resample_shape(patch_shape, factor, n_channels=None):

    """Compute the shape of a resampled array given a shape factor.

    :param patch_shape: size of the initial array (without number of channels).

    :param factor: resampling factor. Can be a number, sequence, or 1d numpy array.

    :param n_channels: (optional) if not None, add a number of channel at the end of the computed shape.

    :return: list containing the shape of the input array after being resampled by the given factor.

    """

    factor = reformat_to_list(factor, length=len(patch_shape))

    shape = [math.ceil(patch_shape[i] * factor[i]) for i in range(len(patch_shape))]

    if n_channels is not None:

        shape += [n_channels]

    return shape

def add_axis(x, axis=0):

    """Add axis to a numpy array.

    :param x: input array

    :param axis: index of the new axis to add. Can also be a list of indices to add several axes at the same time."""

    axis = reformat_to_list(axis)

    for ax in axis:

        x = np.expand_dims(x, axis=ax)

    return x

def get_padding_margin(cropping, loss_cropping):

    """Compute padding margin"""

    if (cropping is not None) & (loss_cropping is not None):

        cropping = reformat_to_list(cropping)

        loss_cropping = reformat_to_list(loss_cropping)

        n_dims = max(len(cropping), len(loss_cropping))

        cropping = reformat_to_list(cropping, length=n_dims)

        loss_cropping = reformat_to_list(loss_cropping, length=n_dims)

        padding_margin = [int((cropping[i] - loss_cropping[i]) / 2) for i in range(n_dims)]

        if len(padding_margin) == 1:

            padding_margin = padding_margin[0]

    else:

        padding_margin = None

    return padding_margin

# -------------------------------------------- build affine matrices/tensors -------------------------------------------

def create_affine_transformation_matrix(n_dims, scaling=None, rotation=None, shearing=None, translation=None):

    """Create a 4x4 affine transformation matrix from specified values

    :param n_dims: integer, can either be 2 or 3.

    :param scaling: list of 3 scaling values

    :param rotation: list of 3 angles (degrees) for rotations around 1st, 2nd, 3rd axis

    :param shearing: list of 6 shearing values

    :param translation: list of 3 values

    :return: 4x4 numpy matrix

    """

    T_scaling = np.eye(n_dims + 1)

    T_shearing = np.eye(n_dims + 1)

    T_translation = np.eye(n_dims + 1)

    if scaling is not None:

        T_scaling[np.arange(n_dims + 1), np.arange(n_dims + 1)] = np.append(scaling, 1)

    if shearing is not None:

        shearing_index = np.ones((n_dims + 1, n_dims + 1), dtype='bool')

        shearing_index[np.eye(n_dims + 1, dtype='bool')] = False

        shearing_index[-1, :] = np.zeros((n_dims + 1))

        shearing_index[:, -1] = np.zeros((n_dims + 1))

        T_shearing[shearing_index] = shearing

    if translation is not None:

        T_translation[np.arange(n_dims), n_dims * np.ones(n_dims, dtype='int')] = translation

    if n_dims == 2:

        if rotation is None:

            rotation = np.zeros(1)

        else:

            rotation = np.asarray(rotation) * (math.pi / 180)

        T_rot = np.eye(n_dims + 1)

        T_rot[np.array([0, 1, 0, 1]), np.array([0, 0, 1, 1])] = [np.cos(rotation[0]), np.sin(rotation[0]),

                                                                 np.sin(rotation[0]) * -1, np.cos(rotation[0])]

        return T_translation @ T_rot @ T_shearing @ T_scaling

    else:

        if rotation is None:

            rotation = np.zeros(n_dims)

        else:

            rotation = np.asarray(rotation) * (math.pi / 180)

        T_rot1 = np.eye(n_dims + 1)

        T_rot1[np.array([1, 2, 1, 2]), np.array([1, 1, 2, 2])] = [np.cos(rotation[0]), np.sin(rotation[0]),

                                                                  np.sin(rotation[0]) * -1, np.cos(rotation[0])]

        T_rot2 = np.eye(n_dims + 1)

        T_rot2[np.array([0, 2, 0, 2]), np.array([0, 0, 2, 2])] = [np.cos(rotation[1]), np.sin(rotation[1]) * -1,

                                                                  np.sin(rotation[1]), np.cos(rotation[1])]

        T_rot3 = np.eye(n_dims + 1)

        T_rot3[np.array([0, 1, 0, 1]), np.array([0, 0, 1, 1])] = [np.cos(rotation[2]), np.sin(rotation[2]),

                                                                  np.sin(rotation[2]) * -1, np.cos(rotation[2])]

        return T_translation @ T_rot3 @ T_rot2 @ T_rot1 @ T_shearing @ T_scaling

def sample_affine_transform(batchsize,

                            n_dims,

                            rotation_bounds=False,

                            scaling_bounds=False,

                            shearing_bounds=False,

                            translation_bounds=False,

                            enable_90_rotations=False):

    """build batchsize x 4 x 4 tensor representing an affine transformation in homogeneous coordinates.

    If return_inv is True, also returns the inverse of the created affine matrix."""

    if (rotation_bounds is not False) | (enable_90_rotations is not False):

        if n_dims == 2:

            if rotation_bounds is not False:

                rotation = draw_value_from_distribution(rotation_bounds,

                                                        size=1,

                                                        default_range=15.0,

                                                        return_as_tensor=True,

                                                        batchsize=batchsize)

            else:

                rotation = tf.zeros(tf.concat([batchsize, tf.ones(1, dtype='int32')], axis=0))

        else:  # n_dims = 3

            if rotation_bounds is not False:

                rotation = draw_value_from_distribution(rotation_bounds,

                                                        size=n_dims,

                                                        default_range=15.0,

                                                        return_as_tensor=True,

                                                        batchsize=batchsize)

            else:

                rotation = tf.zeros(tf.concat([batchsize, 3 * tf.ones(1, dtype='int32')], axis=0))

        if enable_90_rotations:

            rotation = tf.cast(tf.random.uniform(tf.shape(rotation), maxval=4, dtype='int32') * 90, 'float32') \

                       + rotation

        T_rot = create_rotation_transform(rotation, n_dims)

    else:

        T_rot = tf.tile(tf.expand_dims(tf.eye(n_dims), axis=0),

                        tf.concat([batchsize, tf.ones(2, dtype='int32')], axis=0))

    if shearing_bounds is not False:

        shearing = draw_value_from_distribution(shearing_bounds,

                                                size=n_dims ** 2 - n_dims,

                                                default_range=.01,

                                                return_as_tensor=True,

                                                batchsize=batchsize)

        T_shearing = create_shearing_transform(shearing, n_dims)

    else:

        T_shearing = tf.tile(tf.expand_dims(tf.eye(n_dims), axis=0),

                             tf.concat([batchsize, tf.ones(2, dtype='int32')], axis=0))

    if scaling_bounds is not False:

        scaling = draw_value_from_distribution(scaling_bounds,

                                               size=n_dims,

                                               centre=1,

                                               default_range=.15,

                                               return_as_tensor=True,

                                               batchsize=batchsize)

        T_scaling = tf.linalg.diag(scaling)

    else:

        T_scaling = tf.tile(tf.expand_dims(tf.eye(n_dims), axis=0),

                            tf.concat([batchsize, tf.ones(2, dtype='int32')], axis=0))

    T = tf.matmul(T_scaling, tf.matmul(T_shearing, T_rot))

    if translation_bounds is not False:

        translation = draw_value_from_distribution(translation_bounds,

                                                   size=n_dims,

                                                   default_range=5,

                                                   return_as_tensor=True,

                                                   batchsize=batchsize)

        T = tf.concat([T, tf.expand_dims(translation, axis=-1)], axis=-1)

    else:

        T = tf.concat([T, tf.zeros(tf.concat([tf.shape(T)[:2], tf.ones(1, dtype='int32')], 0))], axis=-1)

    # build rigid transform

    T_last_row = tf.expand_dims(tf.concat([tf.zeros((1, n_dims)), tf.ones((1, 1))], axis=1), 0)

    T_last_row = tf.tile(T_last_row, tf.concat([batchsize, tf.ones(2, dtype='int32')], axis=0))

    T = tf.concat([T, T_last_row], axis=1)

    return T

def create_rotation_transform(rotation, n_dims):

    """build rotation transform from 3d or 2d rotation coefficients. Angles are given in degrees."""

    rotation = rotation * np.pi / 180

    if n_dims == 3:

        shape = tf.shape(tf.expand_dims(rotation[..., 0], -1))

        Rx_row0 = tf.expand_dims(tf.tile(tf.expand_dims(tf.convert_to_tensor([1., 0., 0.]), 0), shape), axis=1)

        Rx_row1 = tf.stack([tf.zeros(shape), tf.expand_dims(tf.cos(rotation[..., 0]), -1),

                            tf.expand_dims(-tf.sin(rotation[..., 0]), -1)], axis=-1)

        Rx_row2 = tf.stack([tf.zeros(shape), tf.expand_dims(tf.sin(rotation[..., 0]), -1),

                            tf.expand_dims(tf.cos(rotation[..., 0]), -1)], axis=-1)

        Rx = tf.concat([Rx_row0, Rx_row1, Rx_row2], axis=1)

        Ry_row0 = tf.stack([tf.expand_dims(tf.cos(rotation[..., 1]), -1), tf.zeros(shape),

                            tf.expand_dims(tf.sin(rotation[..., 1]), -1)], axis=-1)

        Ry_row1 = tf.expand_dims(tf.tile(tf.expand_dims(tf.convert_to_tensor([0., 1., 0.]), 0), shape), axis=1)

        Ry_row2 = tf.stack([tf.expand_dims(-tf.sin(rotation[..., 1]), -1), tf.zeros(shape),

                            tf.expand_dims(tf.cos(rotation[..., 1]), -1)], axis=-1)

        Ry = tf.concat([Ry_row0, Ry_row1, Ry_row2], axis=1)

        Rz_row0 = tf.stack([tf.expand_dims(tf.cos(rotation[..., 2]), -1),

                            tf.expand_dims(-tf.sin(rotation[..., 2]), -1), tf.zeros(shape)], axis=-1)

        Rz_row1 = tf.stack([tf.expand_dims(tf.sin(rotation[..., 2]), -1),

                            tf.expand_dims(tf.cos(rotation[..., 2]), -1), tf.zeros(shape)], axis=-1)

        Rz_row2 = tf.expand_dims(tf.tile(tf.expand_dims(tf.convert_to_tensor([0., 0., 1.]), 0), shape), axis=1)

        Rz = tf.concat([Rz_row0, Rz_row1, Rz_row2], axis=1)

        T_rot = tf.matmul(tf.matmul(Rx, Ry), Rz)

    elif n_dims == 2:

        R_row0 = tf.stack([tf.expand_dims(tf.cos(rotation[..., 0]), -1),

                           tf.expand_dims(tf.sin(rotation[..., 0]), -1)], axis=-1)

        R_row1 = tf.stack([tf.expand_dims(-tf.sin(rotation[..., 0]), -1),

                           tf.expand_dims(tf.cos(rotation[..., 0]), -1)], axis=-1)

        T_rot = tf.concat([R_row0, R_row1], axis=1)

    else:

        raise Exception('only supports 2 or 3D.')

    return T_rot

def create_shearing_transform(shearing, n_dims):

    """build shearing transform from 2d/3d shearing coefficients"""

    shape = tf.shape(tf.expand_dims(shearing[..., 0], -1))

    if n_dims == 3:

        shearing_row0 = tf.stack([tf.ones(shape), tf.expand_dims(shearing[..., 0], -1),

                                  tf.expand_dims(shearing[..., 1], -1)], axis=-1)

        shearing_row1 = tf.stack([tf.expand_dims(shearing[..., 2], -1), tf.ones(shape),

                                  tf.expand_dims(shearing[..., 3], -1)], axis=-1)

        shearing_row2 = tf.stack([tf.expand_dims(shearing[..., 4], -1), tf.expand_dims(shearing[..., 5], -1),

                                  tf.ones(shape)], axis=-1)

        T_shearing = tf.concat([shearing_row0, shearing_row1, shearing_row2], axis=1)

    elif n_dims == 2:

        shearing_row0 = tf.stack([tf.ones(shape), tf.expand_dims(shearing[..., 0], -1)], axis=-1)

        shearing_row1 = tf.stack([tf.expand_dims(shearing[..., 1], -1), tf.ones(shape)], axis=-1)

        T_shearing = tf.concat([shearing_row0, shearing_row1], axis=1)

    else:

        raise Exception('only supports 2 or 3D.')

    return T_shearing

# --------------------------------------------------- miscellaneous ----------------------------------------------------

def infer(x):

    """ Try to parse input to float. If it fails, tries boolean, and otherwise keep it as string """

    try:

        x = float(x)

    except ValueError:

        if x == 'False':

            x = False

        elif x == 'True':

            x = True

        elif not isinstance(x, str):

            raise TypeError('input should be an int/float/boolean/str, had {}'.format(type(x)))

    return x

class LoopInfo:

    """

    Class to print the current iteration in a for loop, and optionally the estimated remaining time.

    Instantiate just before the loop, and call the update method at the start of the loop.

    The printed text has the following format:

    processing i/total    remaining time: hh:mm:ss

    """

    def __init__(self, n_iterations, spacing=10, text='processing', print_time=False):

        """

        :param n_iterations: total number of iterations of the for loop.

        :param spacing: frequency at which the update info will be printed on screen.

        :param text: text to print. Default is processing.

        :param print_time: whether to print the estimated remaining time. Default is False.

        """

        # loop parameters

        self.n_iterations = n_iterations

        self.spacing = spacing

        # text parameters

        self.text = text

        self.print_time = print_time

        self.print_previous_time = False

        self.align = len(str(self.n_iterations)) * 2 + 1 + 3

        # timing parameters

        self.iteration_durations = np.zeros((n_iterations,))

        self.start = time.time()

        self.previous = time.time()

    def update(self, idx):

        # time iteration

        now = time.time()

        self.iteration_durations[idx] = now - self.previous

        self.previous = now

        # print text

        if idx == 0:

            print(self.text + ' 1/{}'.format(self.n_iterations))

        elif idx % self.spacing == self.spacing - 1:

            iteration = str(idx + 1) + '/' + str(self.n_iterations)

            if self.print_time:

                # estimate remaining time

                max_duration = np.max(self.iteration_durations)

                average_duration = np.mean(self.iteration_durations[self.iteration_durations > .01 * max_duration])

                remaining_time = int(average_duration * (self.n_iterations - idx))

                # print total remaining time only if it is greater than 1s or if it was previously printed

                if (remaining_time > 1) | self.print_previous_time:

                    eta = str(timedelta(seconds=remaining_time))

                    print(self.text + ' {:<{x}} remaining time: {}'.format(iteration, eta, x=self.align))

                    self.print_previous_time = True

                else:

                    print(self.text + ' {}'.format(iteration))

            else:

                print(self.text + ' {}'.format(iteration))

def get_mapping_lut(source, dest=None):

    """This functions returns the look-up table to map a list of N values (source) to another list (dest).

    If the second list is not given, we assume it is equal to [0, ..., N-1]."""

    # initialise

    source = np.array(reformat_to_list(source), dtype='int32')

    n_labels = source.shape[0]

    # build new label list if necessary

    if dest is None:

        dest = np.arange(n_labels, dtype='int32')

    else:

        assert len(source) == len(dest), 'label_list and new_label_list should have the same length'

        dest = np.array(reformat_to_list(dest, dtype='int'))

    # build look-up table

    lut = np.zeros(np.max(source) + 1, dtype='int32')

    for source, dest in zip(source, dest):

        lut[source] = dest

    return lut

def build_training_generator(gen, batchsize):

    """Build generator for training a network."""

    while True:

        inputs = next(gen)

        if batchsize > 1:

            target = np.concatenate([np.zeros((1, 1))] * batchsize, 0)

        else:

            target = np.zeros((1, 1))

        yield inputs, target

def find_closest_number_divisible_by_m(n, m, answer_type='lower'):

    """Return the closest integer to n that is divisible by m. answer_type can either be 'closer', 'lower' (only returns

    values lower than n), or 'higher' (only returns values higher than m)."""

    if n % m == 0:

        return n

    else:

        q = int(n / m)

        lower = q * m

        higher = (q + 1) * m

        if answer_type == 'lower':

            return lower

        elif answer_type == 'higher':

            return higher

        elif answer_type == 'closer':

            return lower if (n - lower) < (higher - n) else higher

        else:

            raise Exception('answer_type should be lower, higher, or closer, had : %s' % answer_type)

def build_binary_structure(connectivity, n_dims, shape=None):

    """Return a dilation/erosion element with provided connectivity"""

    if shape is None:

        shape = [connectivity * 2 + 1] * n_dims

    else:

        shape = reformat_to_list(shape, length=n_dims)

    dist = np.ones(shape)

    center = tuple([tuple([int(s / 2)]) for s in shape])

    dist[center] = 0

    dist = distance_transform_edt(dist)

    struct = (dist <= connectivity) * 1

    return struct

def draw_value_from_distribution(hyperparameter,

                                 size=1,

                                 distribution='uniform',

                                 centre=0.,

                                 default_range=10.0,

                                 positive_only=False,

                                 return_as_tensor=False,

                                 batchsize=None):

    """Sample values from a uniform, or normal distribution of given hyperparameters.

    These hyperparameters are to the number of 2 in both uniform and normal cases.

    :param hyperparameter: values of the hyperparameters. Can either be:

    1) None, in each case the two hyperparameters are given by [center-default_range, center+default_range],

    2) a number, where the two hyperparameters are given by [centre-hyperparameter, centre+hyperparameter],

    3) a sequence of length 2, directly defining the two hyperparameters: [min, max] if the distribution is uniform,

    [mean, std] if the distribution is normal.

    4) a numpy array, with size (2, m). In this case, the function returns a 1d array of size m, where each value has

    been sampled independently with the specified hyperparameters. If the distribution is uniform, rows correspond to

    its lower and upper bounds, and if the distribution is normal, rows correspond to its mean and std deviation.

    5) a numpy array of size (2*n, m). Same as 4) but we first randomly select a block of two rows among the

    n possibilities.

    6) the path to a numpy array corresponding to case 4 or 5.

    7) False, in which case this function returns None.

    :param size: (optional) number of values to sample. All values are sampled independently.

    Used only if hyperparameter is not a numpy array.

    :param distribution: (optional) the distribution type. Can be 'uniform' or 'normal'. Default is 'uniform'.

    :param centre: (optional) default centre to use if hyperparameter is None or a number.

    :param default_range: (optional) default range to use if hyperparameter is None.

    :param positive_only: (optional) whether to reset all negative values to zero.

    :param return_as_tensor: (optional) whether to return the result as a tensorflow tensor

    :param batchsize: (optional) if return_as_tensor is true, then you can sample a tensor of a given batchsize. Give

    this batchsize as a tensorflow tensor here.

    :return: a float, or a numpy 1d array if size > 1, or hyperparameter is itself a numpy array.

    Returns None if hyperparameter is False.

    """

    # return False is hyperparameter is False

    if hyperparameter is False:

        return None

    # reformat parameter_range

    hyperparameter = load_array_if_path(hyperparameter, load_as_numpy=True)

    if not isinstance(hyperparameter, np.ndarray):

        if hyperparameter is None:

            hyperparameter = np.array([[centre - default_range] * size, [centre + default_range] * size])

        elif isinstance(hyperparameter, (int, float)):

            hyperparameter = np.array([[centre - hyperparameter] * size, [centre + hyperparameter] * size])

        elif isinstance(hyperparameter, (list, tuple)):

            assert len(hyperparameter) == 2, 'if list, parameter_range should be of length 2.'

            hyperparameter = np.transpose(np.tile(np.array(hyperparameter), (size, 1)))

        else:

            raise ValueError('parameter_range should either be None, a number, a sequence, or a numpy array.')

    elif isinstance(hyperparameter, np.ndarray):

        assert hyperparameter.shape[0] % 2 == 0, 'number of rows of parameter_range should be divisible by 2'

        n_modalities = int(hyperparameter.shape[0] / 2)

        modality_idx = 2 * np.random.randint(n_modalities)

        hyperparameter = hyperparameter[modality_idx: modality_idx + 2, :]

    # draw values as tensor

    if return_as_tensor:

        shape = KL.Lambda(lambda x: tf.convert_to_tensor(hyperparameter.shape[1], 'int32'))([])

        if batchsize is not None:

            shape = KL.Lambda(lambda x: tf.concat([x[0], tf.expand_dims(x[1], axis=0)], axis=0))([batchsize, shape])

        if distribution == 'uniform':

            parameter_value = KL.Lambda(lambda x: tf.random.uniform(shape=x,

                                                                    minval=hyperparameter[0, :],

                                                                    maxval=hyperparameter[1, :]))(shape)

        elif distribution == 'normal':

            parameter_value = KL.Lambda(lambda x: tf.random.normal(shape=x,

                                                                   mean=hyperparameter[0, :],

                                                                   stddev=hyperparameter[1, :]))(shape)

        else:

            raise ValueError("Distribution not supported, should be 'uniform' or 'normal'.")

        if positive_only:

            parameter_value = KL.Lambda(lambda x: K.clip(x, 0, None))(parameter_value)

    # draw values as numpy array

    else:

        if distribution == 'uniform':

            parameter_value = np.random.uniform(low=hyperparameter[0, :], high=hyperparameter[1, :])

        elif distribution == 'normal':

            parameter_value = np.random.normal(loc=hyperparameter[0, :], scale=hyperparameter[1, :])

        else:

            raise ValueError("Distribution not supported, should be 'uniform' or 'normal'.")

        if positive_only:

            parameter_value[parameter_value < 0] = 0

    return parameter_value

def build_exp(x, first, last, fix_point):

    # first = f(0), last = f(+inf), fix_point = [x0, f(x0))]

    a = last

    b = first - last

    c = - (1 / fix_point[0]) * np.log((fix_point[1] - last) / (first - last))

    return a + b * np.exp(-c * x)