"""

If you use this code, please cite one of the SynthSeg papers:

https://github.com/BBillot/SynthSeg/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 os

import numpy as np

import tensorflow as tf

from keras import models

from keras import layers as KL

# project imports

from SynthSeg import metrics_model as metrics

from SynthSeg.training import train_model

from SynthSeg.labels_to_image_model import get_shapes

from SynthSeg.training_supervised import build_model_inputs

# third-party imports

from ext.lab2im import utils, layers

from ext.neuron import models as nrn_models

def training(list_paths_input_labels,

             list_paths_target_labels,

             model_dir,

             input_segmentation_labels,

             target_segmentation_labels=None,

             subjects_prob=None,

             batchsize=1,

             output_shape=None,

             scaling_bounds=.2,

             rotation_bounds=15,

             shearing_bounds=.012,

             nonlin_std=3.,

             nonlin_scale=.04,

             prob_erosion_dilation=0.3,

             min_erosion_dilation=4,

             max_erosion_dilation=5,

             n_levels=5,

             nb_conv_per_level=2,

             conv_size=5,

             unet_feat_count=16,

             feat_multiplier=2,

             activation='elu',

             skip_n_concatenations=2,

             lr=1e-4,

             wl2_epochs=1,

             dice_epochs=50,

             steps_per_epoch=10000,

             checkpoint=None):

    """

    This function trains a UNet to segment MRI images with synthetic scans generated by sampling a GMM conditioned on

    label maps. We regroup the parameters in four categories: General, Augmentation, Architecture, Training.

    # IMPORTANT !!!

    # Each time we provide a parameter with separate values for each axis (e.g. with a numpy array or a sequence),

    # these values refer to the RAS axes.

    :param list_paths_input_labels: list of all the paths of the input label maps. These correspond to "noisy"

    segmentations that the denoiser will be trained to correct.

    :param list_paths_target_labels: list of all the paths of the output label maps. Must have the same order as

    list_paths_input_labels. These are the target label maps that the network will learn to produce given the "noisy"

    input label maps.

    :param model_dir: path of a directory where the models will be saved during training.

    :param input_segmentation_labels: list of all the label values present in the input label maps.

    :param target_segmentation_labels: list of all the label values present in the output label maps. By default (None)

    this will be taken to be the same as input_segmentation_labels.

    # ----------------------------------------------- General parameters -----------------------------------------------

    # label maps parameters

    :param subjects_prob: (optional) relative order of importance (doesn't have to be probabilistic), with which to pick

    the provided label maps at each minibatch. Can be a sequence, a 1D numpy array, or the path to such an array, and it

    must be as long as path_label_maps. By default, all label maps are chosen with the same importance.

    # output-related parameters

    :param batchsize: (optional) number of images to generate per mini-batch. Default is 1.

    :param output_shape: (optional) desired shape of the output image, obtained by randomly cropping the generated image

    Can be an integer (same size in all dimensions), a sequence, a 1d numpy array, or the path to a 1d numpy array.

    Default is None, where no cropping is performed.

    # --------------------------------------------- Augmentation parameters --------------------------------------------

    # spatial deformation parameters

    :param scaling_bounds: (optional) if apply_linear_trans is True, 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) the path to a numpy array of shape (2, n_dims), in which case the scaling factor in dimension i is sampled from

    the uniform distribution of bounds (scaling_bounds[0, i], scaling_bounds[1, i]) for the i-th dimension.

    3) False, in which case scaling is completely turned off.

    Default is scaling_bounds = 0.2 (case 1)

    :param rotation_bounds: (optional) same as scaling bounds but for the rotation angle, except that for case 1 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 nonlin_std: (optional) Standard deviation of the normal distribution from which we sample the first

    tensor for synthesising the deformation field. Set to 0 to completely deactivate elastic deformation.

    :param nonlin_scale: (optional) Ratio between the size of the input label maps and the size of the sampled

    tensor for synthesising the elastic deformation field.

    # degradation of the input labels

    :param prob_erosion_dilation: (optional) probability with which to degrade the input label maps with erosion or 

    dilation. If 0, then no erosion/dilation is applied to the label maps given as inputs to the network.

    :param min_erosion_dilation: (optional) when prob_erosion_dilation is not zero, erosion and dilation of random

    coefficients are applied. Set the minimum erosion/dilation coefficient here.

    :param max_erosion_dilation: (optional) Set the maximum erosion/dilation coefficient here.

    # ------------------------------------------ UNet architecture parameters ------------------------------------------

    :param n_levels: (optional) number of level for the Unet. Default is 5.

    :param nb_conv_per_level: (optional) number of convolutional layers per level. Default is 2.

    :param conv_size: (optional) size of the convolution kernels. Default is 2.

    :param unet_feat_count: (optional) number of feature for the first layer of the UNet. Default is 24.

    :param feat_multiplier: (optional) multiply the number of feature by this number at each new level. Default is 2.

    :param activation: (optional) activation function. Can be 'elu', 'relu'.

    :param skip_n_concatenations: (optional) number of levels for which to remove the traditional skip connections of

    the UNet architecture. default is zero, which corresponds to the classic UNet architecture. Example:

    If skip_n_concatenations = 2, then we will remove the concatenation link between the two top levels of the UNet.

    # ----------------------------------------------- Training parameters ----------------------------------------------

    :param lr: (optional) learning rate for the training. Default is 1e-4

    :param wl2_epochs: (optional) number of epochs for which the network (except the soft-max layer) is trained with L2

    norm loss function. Default is 1.

    :param dice_epochs: (optional) number of epochs with the soft Dice loss function. Default is 50.

    :param steps_per_epoch: (optional) number of steps per epoch. Default is 10000. Since no online validation is

    possible, this is equivalent to the frequency at which the models are saved.

    :param checkpoint: (optional) path of an already saved model to load before starting the training.

    """

    # check epochs

    assert (wl2_epochs > 0) | (dice_epochs > 0), \

        'either wl2_epochs or dice_epochs must be positive, had {0} and {1}'.format(wl2_epochs, dice_epochs)

    # prepare data files

    input_label_list, _ = utils.get_list_labels(label_list=input_segmentation_labels)

    if target_segmentation_labels is None:

        target_label_list = input_label_list

    else:

        target_label_list, _ = utils.get_list_labels(label_list=target_segmentation_labels)

    n_labels = np.size(target_label_list)

    # create augmentation model

    labels_shape, _, _, _, _, _ = utils.get_volume_info(list_paths_input_labels[0], aff_ref=np.eye(4))

    augmentation_model = build_augmentation_model(labels_shape,

                                                  input_label_list,

                                                  crop_shape=output_shape,

                                                  output_div_by_n=2 ** n_levels,

                                                  scaling_bounds=scaling_bounds,

                                                  rotation_bounds=rotation_bounds,

                                                  shearing_bounds=shearing_bounds,

                                                  nonlin_std=nonlin_std,

                                                  nonlin_scale=nonlin_scale,

                                                  prob_erosion_dilation=prob_erosion_dilation,

                                                  min_erosion_dilation=min_erosion_dilation,

                                                  max_erosion_dilation=max_erosion_dilation)

    unet_input_shape = augmentation_model.output[0].get_shape().as_list()[1:]

    # prepare the segmentation model

    l2l_model = nrn_models.unet(input_model=augmentation_model,

                                input_shape=unet_input_shape,

                                nb_labels=n_labels,

                                nb_levels=n_levels,

                                nb_conv_per_level=nb_conv_per_level,

                                conv_size=conv_size,

                                nb_features=unet_feat_count,

                                feat_mult=feat_multiplier,

                                activation=activation,

                                batch_norm=-1,

                                skip_n_concatenations=skip_n_concatenations,

                                name='l2l')

    # input generator

    model_inputs = build_model_inputs(path_inputs=list_paths_input_labels,

                                      path_outputs=list_paths_target_labels,

                                      batchsize=batchsize,

                                      subjects_prob=subjects_prob,

                                      dtype_input='int32')

    input_generator = utils.build_training_generator(model_inputs, batchsize)

    # pre-training with weighted L2, input is fit to the softmax rather than the probabilities

    if wl2_epochs > 0:

        wl2_model = models.Model(l2l_model.inputs, [l2l_model.get_layer('l2l_likelihood').output])

        wl2_model = metrics.metrics_model(wl2_model, target_label_list, 'wl2')

        train_model(wl2_model, input_generator, lr, wl2_epochs, steps_per_epoch, model_dir, 'wl2', checkpoint)

        checkpoint = os.path.join(model_dir, 'wl2_%03d.h5' % wl2_epochs)

    # fine-tuning with dice metric

    dice_model = metrics.metrics_model(l2l_model, target_label_list, 'dice')

    train_model(dice_model, input_generator, lr, dice_epochs, steps_per_epoch, model_dir, 'dice', checkpoint)

def build_augmentation_model(labels_shape,

                             segmentation_labels,

                             crop_shape=None,

                             output_div_by_n=None,

                             scaling_bounds=0.15,

                             rotation_bounds=15,

                             shearing_bounds=0.012,

                             translation_bounds=False,

                             nonlin_std=3.,

                             nonlin_scale=.0625,

                             prob_erosion_dilation=0.3,

                             min_erosion_dilation=4,

                             max_erosion_dilation=7):

    # reformat resolutions and get shapes

    labels_shape = utils.reformat_to_list(labels_shape)

    n_dims, _ = utils.get_dims(labels_shape)

    n_labels = len(segmentation_labels)

    # get shapes

    crop_shape, _ = get_shapes(labels_shape, crop_shape, np.array([1]*n_dims), np.array([1]*n_dims), output_div_by_n)

    # define model inputs

    net_input = KL.Input(shape=labels_shape + [1], name='l2l_noisy_labels_input', dtype='int32')

    target_input = KL.Input(shape=labels_shape + [1], name='l2l_target_input', dtype='int32')

    # deform labels

    noisy_labels, target = layers.RandomSpatialDeformation(scaling_bounds=scaling_bounds,

                                                           rotation_bounds=rotation_bounds,

                                                           shearing_bounds=shearing_bounds,

                                                           translation_bounds=translation_bounds,

                                                           nonlin_std=nonlin_std,

                                                           nonlin_scale=nonlin_scale,

                                                           inter_method='nearest')([net_input, target_input])

    # cropping

    if crop_shape != labels_shape:

        noisy_labels, target = layers.RandomCrop(crop_shape)([noisy_labels, target])

    # random erosion

    if prob_erosion_dilation > 0:

        noisy_labels = layers.RandomDilationErosion(min_erosion_dilation,

                                                    max_erosion_dilation,

                                                    prob=prob_erosion_dilation)(noisy_labels)

    # convert input labels (i.e. noisy_labels) to [0, ... N-1] and make them one-hot

    noisy_labels = layers.ConvertLabels(np.unique(segmentation_labels))(noisy_labels)

    target = KL.Lambda(lambda x: tf.cast(x[..., 0], 'int32'), name='labels_out')(target)

    noisy_labels = KL.Lambda(lambda x: tf.one_hot(x[0][..., 0], depth=n_labels),

                             name='noisy_labels_out')([noisy_labels, target])

    # build model and return

    brain_model = models.Model(inputs=[net_input, target_input], outputs=[noisy_labels, target])

    return brain_model