"""

tensorflow/keras utilities for the neuron project

If you use this code, please cite 

Dalca AV, Guttag J, Sabuncu MR

Anatomical Priors in Convolutional Networks for Unsupervised Biomedical Segmentation, 

CVPR 2018

Contact: adalca [at] csail [dot] mit [dot] edu

License: GPLv3

"""

import sys

from ext.neuron import layers

# third party

import numpy as np

import tensorflow as tf

import keras

import keras.layers as KL

from keras.models import Model

import keras.backend as K

def unet(nb_features,

         input_shape,

         nb_levels,

         conv_size,

         nb_labels,

         name='unet',

         prefix=None,

         feat_mult=1,

         pool_size=2,

         use_logp=True,

         padding='same',

         dilation_rate_mult=1,

         activation='elu',

         skip_n_concatenations=0,

         use_residuals=False,

         final_pred_activation='softmax',

         nb_conv_per_level=1,

         add_prior_layer=False,

         layer_nb_feats=None,

         conv_dropout=0,

         batch_norm=None,

         input_model=None):

    """

    unet-style keras model with an overdose of parametrization.

    Parameters:

        nb_features: the number of features at each convolutional level

            see below for `feat_mult` and `layer_nb_feats` for modifiers to this number

        input_shape: input layer shape, vector of size ndims + 1 (nb_channels)

        conv_size: the convolution kernel size

        nb_levels: the number of Unet levels (number of downsamples) in the "encoder" 

            (e.g. 4 would give you 4 levels in encoder, 4 in decoder)

        nb_labels: number of output channels

        name (default: 'unet'): the name of the network

        prefix (default: `name` value): prefix to be added to layer names

        feat_mult (default: 1) multiple for `nb_features` as we go down the encoder levels.

            e.g. feat_mult of 2 and nb_features of 16 would yield 32 features in the 

            second layer, 64 features in the third layer, etc.

        pool_size (default: 2): max pooling size (integer or list if specifying per dimension)

        skip_n_concatenations=0: enabled to skip concatenation links between contracting and expanding paths for the n

            top levels.

        use_logp:

        padding:

        dilation_rate_mult:

        activation:

        use_residuals:

        final_pred_activation:

        nb_conv_per_level:

        add_prior_layer:

        skip_n_concatenations:

        layer_nb_feats: list of the number of features for each layer. Automatically used if specified

        conv_dropout: dropout probability

        batch_norm:

        input_model: concatenate the provided input_model to this current model.

            Only the first output of input_model is used.

    """

    # naming

    model_name = name

    if prefix is None:

        prefix = model_name

    # volume size data

    ndims = len(input_shape) - 1

    if isinstance(pool_size, int):

        pool_size = (pool_size,) * ndims

    # get encoding model

    enc_model = conv_enc(nb_features,

                         input_shape,

                         nb_levels,

                         conv_size,

                         name=model_name,

                         prefix=prefix,

                         feat_mult=feat_mult,

                         pool_size=pool_size,

                         padding=padding,

                         dilation_rate_mult=dilation_rate_mult,

                         activation=activation,

                         use_residuals=use_residuals,

                         nb_conv_per_level=nb_conv_per_level,

                         layer_nb_feats=layer_nb_feats,

                         conv_dropout=conv_dropout,

                         batch_norm=batch_norm,

                         input_model=input_model)

    # get decoder

    # use_skip_connections=True makes it a u-net

    lnf = layer_nb_feats[(nb_levels * nb_conv_per_level):] if layer_nb_feats is not None else None

    dec_model = conv_dec(nb_features,

                         [],

                         nb_levels,

                         conv_size,

                         nb_labels,

                         name=model_name,

                         prefix=prefix,

                         feat_mult=feat_mult,

                         pool_size=pool_size,

                         use_skip_connections=True,

                         skip_n_concatenations=skip_n_concatenations,

                         padding=padding,

                         dilation_rate_mult=dilation_rate_mult,

                         activation=activation,

                         use_residuals=use_residuals,

                         final_pred_activation='linear' if add_prior_layer else final_pred_activation,

                         nb_conv_per_level=nb_conv_per_level,

                         batch_norm=batch_norm,

                         layer_nb_feats=lnf,

                         conv_dropout=conv_dropout,

                         input_model=enc_model)

    final_model = dec_model

    if add_prior_layer:

        final_model = add_prior(dec_model,

                                [*input_shape[:-1], nb_labels],

                                name=model_name + '_prior',

                                use_logp=use_logp,

                                final_pred_activation=final_pred_activation)

    return final_model

def ae(nb_features,

       input_shape,

       nb_levels,

       conv_size,

       nb_labels,

       enc_size,

       name='ae',

       feat_mult=1,

       pool_size=2,

       padding='same',

       activation='elu',

       use_residuals=False,

       nb_conv_per_level=1,

       batch_norm=None,

       enc_batch_norm=None,

       ae_type='conv',  # 'dense', or 'conv'

       enc_lambda_layers=None,

       add_prior_layer=False,

       use_logp=True,

       conv_dropout=0,

       include_mu_shift_layer=False,

       single_model=False,  # whether to return a single model, or a tuple of models that can be stacked.

       final_pred_activation='softmax',

       do_vae=False,

       input_model=None):

    """Convolutional Auto-Encoder. Optionally Variational (if do_vae is set to True)."""

    # naming

    model_name = name

    # volume size data

    ndims = len(input_shape) - 1

    if isinstance(pool_size, int):

        pool_size = (pool_size,) * ndims

    # get encoding model

    enc_model = conv_enc(nb_features,

                         input_shape,

                         nb_levels,

                         conv_size,

                         name=model_name,

                         feat_mult=feat_mult,

                         pool_size=pool_size,

                         padding=padding,

                         activation=activation,

                         use_residuals=use_residuals,

                         nb_conv_per_level=nb_conv_per_level,

                         conv_dropout=conv_dropout,

                         batch_norm=batch_norm,

                         input_model=input_model)

    # middle AE structure

    if single_model:

        in_input_shape = None

        in_model = enc_model

    else:

        in_input_shape = enc_model.output.shape.as_list()[1:]

        in_model = None

    mid_ae_model = single_ae(enc_size,

                             in_input_shape,

                             conv_size=conv_size,

                             name=model_name,

                             ae_type=ae_type,

                             input_model=in_model,

                             batch_norm=enc_batch_norm,

                             enc_lambda_layers=enc_lambda_layers,

                             include_mu_shift_layer=include_mu_shift_layer,

                             do_vae=do_vae)

    # decoder

    if single_model:

        in_input_shape = None

        in_model = mid_ae_model

    else:

        in_input_shape = mid_ae_model.output.shape.as_list()[1:]

        in_model = None

    dec_model = conv_dec(nb_features,

                         in_input_shape,

                         nb_levels,

                         conv_size,

                         nb_labels,

                         name=model_name,

                         feat_mult=feat_mult,

                         pool_size=pool_size,

                         use_skip_connections=False,

                         padding=padding,

                         activation=activation,

                         use_residuals=use_residuals,

                         final_pred_activation='linear',

                         nb_conv_per_level=nb_conv_per_level,

                         batch_norm=batch_norm,

                         conv_dropout=conv_dropout,

                         input_model=in_model)

    if add_prior_layer:

        dec_model = add_prior(dec_model,

                              [*input_shape[:-1], nb_labels],

                              name=model_name,

                              prefix=model_name + '_prior',

                              use_logp=use_logp,

                              final_pred_activation=final_pred_activation)

    if single_model:

        return dec_model

    else:

        return dec_model, mid_ae_model, enc_model

def conv_enc(nb_features,

             input_shape,

             nb_levels,

             conv_size,

             name=None,

             prefix=None,

             feat_mult=1,

             pool_size=2,

             dilation_rate_mult=1,

             padding='same',

             activation='elu',

             layer_nb_feats=None,

             use_residuals=False,

             nb_conv_per_level=2,

             conv_dropout=0,

             batch_norm=None,

             input_model=None):

    """Fully Convolutional Encoder"""

    # naming

    model_name = name

    if prefix is None:

        prefix = model_name

    # first layer: input

    name = '%s_input' % prefix

    if input_model is None:

        input_tensor = KL.Input(shape=input_shape, name=name)

        last_tensor = input_tensor

    else:

        input_tensor = input_model.inputs

        last_tensor = input_model.outputs

        if isinstance(last_tensor, list):

            last_tensor = last_tensor[0]

    # volume size data

    ndims = len(input_shape) - 1

    if isinstance(pool_size, int):

        pool_size = (pool_size,) * ndims

    # prepare layers

    convL = getattr(KL, 'Conv%dD' % ndims)

    conv_kwargs = {'padding': padding, 'activation': activation, 'data_format': 'channels_last'}

    maxpool = getattr(KL, 'MaxPooling%dD' % ndims)

    # down arm:

    # add nb_levels of conv + ReLu + conv + ReLu. Pool after each of first nb_levels - 1 layers

    lfidx = 0  # level feature index

    for level in range(nb_levels):

        lvl_first_tensor = last_tensor

        nb_lvl_feats = np.round(nb_features * feat_mult ** level).astype(int)

        conv_kwargs['dilation_rate'] = dilation_rate_mult ** level

        for conv in range(nb_conv_per_level):  # does several conv per level, max pooling applied at the end

            if layer_nb_feats is not None:  # None or List of all the feature numbers

                nb_lvl_feats = layer_nb_feats[lfidx]

                lfidx += 1

            name = '%s_conv_downarm_%d_%d' % (prefix, level, conv)

            if conv < (nb_conv_per_level - 1) or (not use_residuals):

                last_tensor = convL(nb_lvl_feats, conv_size, **conv_kwargs, name=name)(last_tensor)

            else:  # no activation

                last_tensor = convL(nb_lvl_feats, conv_size, padding=padding, name=name)(last_tensor)

            if conv_dropout > 0:

                # conv dropout along feature space only

                name = '%s_dropout_downarm_%d_%d' % (prefix, level, conv)

                noise_shape = [None, *[1] * ndims, nb_lvl_feats]

                last_tensor = KL.Dropout(conv_dropout, noise_shape=noise_shape, name=name)(last_tensor)

        if use_residuals:

            convarm_layer = last_tensor

            # the "add" layer is the original input

            # However, it may not have the right number of features to be added

            nb_feats_in = lvl_first_tensor.get_shape()[-1]

            nb_feats_out = convarm_layer.get_shape()[-1]

            add_layer = lvl_first_tensor

            if nb_feats_in > 1 and nb_feats_out > 1 and (nb_feats_in != nb_feats_out):

                name = '%s_expand_down_merge_%d' % (prefix, level)

                last_tensor = convL(nb_lvl_feats, conv_size, **conv_kwargs, name=name)(lvl_first_tensor)

                add_layer = last_tensor

                if conv_dropout > 0:

                    noise_shape = [None, *[1] * ndims, nb_lvl_feats]

                    convarm_layer = KL.Dropout(conv_dropout, noise_shape=noise_shape)(last_tensor)

            name = '%s_res_down_merge_%d' % (prefix, level)

            last_tensor = KL.add([add_layer, convarm_layer], name=name)

            name = '%s_res_down_merge_act_%d' % (prefix, level)

            last_tensor = KL.Activation(activation, name=name)(last_tensor)

        if batch_norm is not None:

            name = '%s_bn_down_%d' % (prefix, level)

            last_tensor = KL.BatchNormalization(axis=batch_norm, name=name)(last_tensor)

        # max pool if we're not at the last level

        if level < (nb_levels - 1):

            name = '%s_maxpool_%d' % (prefix, level)

            last_tensor = maxpool(pool_size=pool_size, name=name, padding=padding)(last_tensor)

    # create the model and return

    model = Model(inputs=input_tensor, outputs=[last_tensor], name=model_name)

    return model

def conv_dec(nb_features,

             input_shape,

             nb_levels,

             conv_size,

             nb_labels,

             name=None,

             prefix=None,

             feat_mult=1,

             pool_size=2,

             use_skip_connections=False,

             skip_n_concatenations=0,

             padding='same',

             dilation_rate_mult=1,

             activation='elu',

             use_residuals=False,

             final_pred_activation='softmax',

             nb_conv_per_level=2,

             layer_nb_feats=None,

             batch_norm=None,

             conv_dropout=0,

             input_model=None):

    """Fully Convolutional Decoder"""

    # naming

    model_name = name

    if prefix is None:

        prefix = model_name

    # if using skip connections, make sure need to use them.

    if use_skip_connections:

        assert input_model is not None, "is using skip connections, tensors dictionary is required"

    # first layer: input

    input_name = '%s_input' % prefix

    if input_model is None:

        input_tensor = KL.Input(shape=input_shape, name=input_name)

        last_tensor = input_tensor

    else:

        input_tensor = input_model.input

        last_tensor = input_model.output

        input_shape = last_tensor.shape.as_list()[1:]

    # vol size info

    ndims = len(input_shape) - 1

    if isinstance(pool_size, int):

        if ndims > 1:

            pool_size = (pool_size,) * ndims

    # prepare layers

    convL = getattr(KL, 'Conv%dD' % ndims)

    conv_kwargs = {'padding': padding, 'activation': activation}

    upsample = getattr(KL, 'UpSampling%dD' % ndims)

    # up arm:

    # nb_levels - 1 layers of Deconvolution3D

    #    (approx via up + conv + ReLu) + merge + conv + ReLu + conv + ReLu

    lfidx = 0

    for level in range(nb_levels - 1):

        nb_lvl_feats = np.round(nb_features * feat_mult ** (nb_levels - 2 - level)).astype(int)

        conv_kwargs['dilation_rate'] = dilation_rate_mult ** (nb_levels - 2 - level)

        # upsample matching the max pooling layers size

        name = '%s_up_%d' % (prefix, nb_levels + level)

        last_tensor = upsample(size=pool_size, name=name)(last_tensor)

        up_tensor = last_tensor

        # merge layers combining previous layer

        if use_skip_connections & (level < (nb_levels - skip_n_concatenations - 1)):

            conv_name = '%s_conv_downarm_%d_%d' % (prefix, nb_levels - 2 - level, nb_conv_per_level - 1)

            cat_tensor = input_model.get_layer(conv_name).output

            name = '%s_merge_%d' % (prefix, nb_levels + level)

            last_tensor = KL.concatenate([cat_tensor, last_tensor], axis=ndims + 1, name=name)

        # convolution layers

        for conv in range(nb_conv_per_level):

            if layer_nb_feats is not None:

                nb_lvl_feats = layer_nb_feats[lfidx]

                lfidx += 1

            name = '%s_conv_uparm_%d_%d' % (prefix, nb_levels + level, conv)

            if conv < (nb_conv_per_level - 1) or (not use_residuals):

                last_tensor = convL(nb_lvl_feats, conv_size, **conv_kwargs, name=name)(last_tensor)

            else:

                last_tensor = convL(nb_lvl_feats, conv_size, padding=padding, name=name)(last_tensor)

            if conv_dropout > 0:

                name = '%s_dropout_uparm_%d_%d' % (prefix, level, conv)

                noise_shape = [None, *[1] * ndims, nb_lvl_feats]

                last_tensor = KL.Dropout(conv_dropout, noise_shape=noise_shape, name=name)(last_tensor)

        # residual block

        if use_residuals:

            # the "add" layer is the original input

            # However, it may not have the right number of features to be added

            add_layer = up_tensor

            nb_feats_in = add_layer.get_shape()[-1]

            nb_feats_out = last_tensor.get_shape()[-1]

            if nb_feats_in > 1 and nb_feats_out > 1 and (nb_feats_in != nb_feats_out):

                name = '%s_expand_up_merge_%d' % (prefix, level)

                add_layer = convL(nb_lvl_feats, conv_size, **conv_kwargs, name=name)(add_layer)

                if conv_dropout > 0:

                    noise_shape = [None, *[1] * ndims, nb_lvl_feats]

                    last_tensor = KL.Dropout(conv_dropout, noise_shape=noise_shape)(last_tensor)

            name = '%s_res_up_merge_%d' % (prefix, level)

            last_tensor = KL.add([last_tensor, add_layer], name=name)

            name = '%s_res_up_merge_act_%d' % (prefix, level)

            last_tensor = KL.Activation(activation, name=name)(last_tensor)

        if batch_norm is not None:

            name = '%s_bn_up_%d' % (prefix, level)

            last_tensor = KL.BatchNormalization(axis=batch_norm, name=name)(last_tensor)

    # Compute likelihood prediction (no activation yet)

    name = '%s_likelihood' % prefix

    last_tensor = convL(nb_labels, 1, activation=None, name=name)(last_tensor)

    like_tensor = last_tensor

    # output prediction layer

    # we use a softmax to compute P(L_x|I) where x is each location

    if final_pred_activation == 'softmax':

        name = '%s_prediction' % prefix

        softmax_lambda_fcn = lambda x: keras.activations.softmax(x, axis=ndims + 1)

        pred_tensor = KL.Lambda(softmax_lambda_fcn, name=name)(last_tensor)

    # otherwise create a layer that does nothing.

    else:

        name = '%s_prediction' % prefix

        pred_tensor = KL.Activation('linear', name=name)(like_tensor)

    # create the model and return

    model = Model(inputs=input_tensor, outputs=pred_tensor, name=model_name)

    return model

def add_prior(input_model,

              prior_shape,

              name='prior_model',

              prefix=None,

              use_logp=True,

              final_pred_activation='softmax'):

    """

    Append post-prior layer to a given model

    """

    # naming

    model_name = name

    if prefix is None:

        prefix = model_name

    # prior input layer

    prior_input_name = '%s-input' % prefix

    prior_tensor = KL.Input(shape=prior_shape, name=prior_input_name)

    prior_tensor_input = prior_tensor

    like_tensor = input_model.output

    # operation varies depending on whether we log() prior or not.

    if use_logp:

        print("Breaking change: use_logp option now requires log input!", file=sys.stderr)

        merge_op = KL.add

    else:

        # using sigmoid to get the likelihood values between 0 and 1

        # note: they won't add up to 1.

        name = '%s_likelihood_sigmoid' % prefix

        like_tensor = KL.Activation('sigmoid', name=name)(like_tensor)

        merge_op = KL.multiply

    # merge the likelihood and prior layers into posterior layer

    name = '%s_posterior' % prefix

    post_tensor = merge_op([prior_tensor, like_tensor], name=name)

    # output prediction layer

    # we use a softmax to compute P(L_x|I) where x is each location

    pred_name = '%s_prediction' % prefix

    if final_pred_activation == 'softmax':

        assert use_logp, 'cannot do softmax when adding prior via P()'

        print("using final_pred_activation %s for %s" % (final_pred_activation, model_name))

        softmax_lambda_fcn = lambda x: keras.activations.softmax(x, axis=-1)

        pred_tensor = KL.Lambda(softmax_lambda_fcn, name=pred_name)(post_tensor)

    else:

        pred_tensor = KL.Activation('linear', name=pred_name)(post_tensor)

    # create the model

    model_inputs = [*input_model.inputs, prior_tensor_input]

    model = Model(inputs=model_inputs, outputs=[pred_tensor], name=model_name)

    # compile

    return model

def single_ae(enc_size,

              input_shape,

              name='single_ae',

              prefix=None,

              ae_type='dense',  # 'dense', or 'conv'

              conv_size=None,

              input_model=None,

              enc_lambda_layers=None,

              batch_norm=True,

              padding='same',

              activation=None,

              include_mu_shift_layer=False,

              do_vae=False):

    """single-layer Autoencoder (i.e. input - encoding - output"""

    # naming

    model_name = name

    if prefix is None:

        prefix = model_name

    if enc_lambda_layers is None:

        enc_lambda_layers = []

    # prepare input

    input_name = '%s_input' % prefix

    if input_model is None:

        assert input_shape is not None, 'input_shape of input_model is necessary'

        input_tensor = KL.Input(shape=input_shape, name=input_name)

        last_tensor = input_tensor

    else:

        input_tensor = input_model.input

        last_tensor = input_model.output

        input_shape = last_tensor.shape.as_list()[1:]

    input_nb_feats = last_tensor.shape.as_list()[-1]

    # prepare conv type based on input

    ndims = len(input_shape) - 1

    if ae_type == 'conv':

        convL = getattr(KL, 'Conv%dD' % ndims)

        assert conv_size is not None, 'with conv ae, need conv_size'

        conv_kwargs = {'padding': padding, 'activation': activation}

        enc_size_str = None

    # if want to go through a dense layer in the middle of the U, need to:

    # - flatten last layer if not flat

    # - do dense encoding and decoding

    # - unflatten (reshape spatially) at end

    else:  # ae_type == 'dense'

        if len(input_shape) > 1:

            name = '%s_ae_%s_down_flat' % (prefix, ae_type)

            last_tensor = KL.Flatten(name=name)(last_tensor)

        convL = conv_kwargs = None

        assert len(enc_size) == 1, "enc_size should be of length 1 for dense layer"

        enc_size_str = ''.join(['%d_' % d for d in enc_size])[:-1]

    # recall this layer

    pre_enc_layer = last_tensor

    # encoding layer

    if ae_type == 'dense':

        name = '%s_ae_mu_enc_dense_%s' % (prefix, enc_size_str)

        last_tensor = KL.Dense(enc_size[0], name=name)(pre_enc_layer)

    else:  # convolution

        # convolve then resize. enc_size should be [nb_dim1, nb_dim2, ..., nb_feats]

        assert len(enc_size) == len(input_shape), \

            "encoding size does not match input shape %d %d" % (len(enc_size), len(input_shape))

        if list(enc_size)[:-1] != list(input_shape)[:-1] and \

                all([f is not None for f in input_shape[:-1]]) and \

                all([f is not None for f in enc_size[:-1]]):

            name = '%s_ae_mu_enc_conv' % prefix

            last_tensor = convL(enc_size[-1], conv_size, name=name, **conv_kwargs)(pre_enc_layer)

            name = '%s_ae_mu_enc' % prefix

            zf = [enc_size[:-1][f] / last_tensor.shape.as_list()[1:-1][f] for f in range(len(enc_size) - 1)]

            last_tensor = layers.Resize(zoom_factor=zf, name=name)(last_tensor)

        elif enc_size[-1] is None:  # convolutional, but won't tell us bottleneck

            name = '%s_ae_mu_enc' % prefix

            last_tensor = KL.Lambda(lambda x: x, name=name)(pre_enc_layer)

        else:

            name = '%s_ae_mu_enc' % prefix

            last_tensor = convL(enc_size[-1], conv_size, name=name, **conv_kwargs)(pre_enc_layer)

    if include_mu_shift_layer:

        # shift

        name = '%s_ae_mu_shift' % prefix

        last_tensor = layers.LocalBias(name=name)(last_tensor)

    # encoding clean-up layers

    for layer_fcn in enc_lambda_layers:

        lambda_name = layer_fcn.__name__

        name = '%s_ae_mu_%s' % (prefix, lambda_name)

        last_tensor = KL.Lambda(layer_fcn, name=name)(last_tensor)

    if batch_norm is not None:

        name = '%s_ae_mu_bn' % prefix

        last_tensor = KL.BatchNormalization(axis=batch_norm, name=name)(last_tensor)

    # have a simple layer that does nothing to have a clear name before sampling

    name = '%s_ae_mu' % prefix

    last_tensor = KL.Lambda(lambda x: x, name=name)(last_tensor)

    # if doing variational AE, will need the sigma layer as well.

    if do_vae:

        mu_tensor = last_tensor

        # encoding layer

        if ae_type == 'dense':

            name = '%s_ae_sigma_enc_dense_%s' % (prefix, enc_size_str)

            last_tensor = KL.Dense(enc_size[0], name=name)(pre_enc_layer)

        else:

            if list(enc_size)[:-1] != list(input_shape)[:-1] and \

                    all([f is not None for f in input_shape[:-1]]) and \

                    all([f is not None for f in enc_size[:-1]]):

                assert len(enc_size) - 1 == 2, "Sorry, I have not yet implemented non-2D resizing..."

                name = '%s_ae_sigma_enc_conv' % prefix

                last_tensor = convL(enc_size[-1], conv_size, name=name, **conv_kwargs)(pre_enc_layer)

                name = '%s_ae_sigma_enc' % prefix

                resize_fn = lambda x: tf.image.resize_bilinear(x, enc_size[:-1])

                last_tensor = KL.Lambda(resize_fn, name=name)(last_tensor)

            elif enc_size[-1] is None:  # convolutional, but won't tell us bottleneck

                name = '%s_ae_sigma_enc' % prefix

                last_tensor = convL(pre_enc_layer.shape.as_list()[-1], conv_size, name=name, **conv_kwargs)(

                    pre_enc_layer)

                # cannot use lambda, then mu and sigma will be same layer.

                # last_tensor = KL.Lambda(lambda x: x, name=name)(pre_enc_layer)

            else:

                name = '%s_ae_sigma_enc' % prefix

                last_tensor = convL(enc_size[-1], conv_size, name=name, **conv_kwargs)(pre_enc_layer)

        # encoding clean-up layers

        for layer_fcn in enc_lambda_layers:

            lambda_name = layer_fcn.__name__

            name = '%s_ae_sigma_%s' % (prefix, lambda_name)

            last_tensor = KL.Lambda(layer_fcn, name=name)(last_tensor)

        if batch_norm is not None:

            name = '%s_ae_sigma_bn' % prefix

            last_tensor = KL.BatchNormalization(axis=batch_norm, name=name)(last_tensor)

        # have a simple layer that does nothing to have a clear name before sampling

        name = '%s_ae_sigma' % prefix

        last_tensor = KL.Lambda(lambda x: x, name=name)(last_tensor)

        logvar_tensor = last_tensor

        # VAE sampling

        sampler = _VAESample().sample_z

        name = '%s_ae_sample' % prefix

        last_tensor = KL.Lambda(sampler, name=name)([mu_tensor, logvar_tensor])

    if include_mu_shift_layer:

        # shift

        name = '%s_ae_sample_shift' % prefix

        last_tensor = layers.LocalBias(name=name)(last_tensor)

    # decoding layer

    if ae_type == 'dense':

        name = '%s_ae_%s_dec_flat_%s' % (prefix, ae_type, enc_size_str)

        last_tensor = KL.Dense(np.prod(input_shape), name=name)(last_tensor)

        # unflatten if dense method

        if len(input_shape) > 1:

            name = '%s_ae_%s_dec' % (prefix, ae_type)

            last_tensor = KL.Reshape(input_shape, name=name)(last_tensor)

    else:

        if list(enc_size)[:-1] != list(input_shape)[:-1] and \

                all([f is not None for f in input_shape[:-1]]) and \

                all([f is not None for f in enc_size[:-1]]):

            name = '%s_ae_mu_dec' % prefix

            zf = [last_tensor.shape.as_list()[1:-1][f] / enc_size[:-1][f] for f in range(len(enc_size) - 1)]

            last_tensor = layers.Resize(zoom_factor=zf, name=name)(last_tensor)

        name = '%s_ae_%s_dec' % (prefix, ae_type)

        last_tensor = convL(input_nb_feats, conv_size, name=name, **conv_kwargs)(last_tensor)

    if batch_norm is not None:

        name = '%s_bn_ae_%s_dec' % (prefix, ae_type)

        last_tensor = KL.BatchNormalization(axis=batch_norm, name=name)(last_tensor)

    # create the model and return

    model = Model(inputs=input_tensor, outputs=[last_tensor], name=model_name)

    return model

###############################################################################

# Helper function

###############################################################################

class _VAESample:

    def __init__(self):

        pass

    def sample_z(self, args):

        mu, log_var = args

        shape = K.shape(mu)

        eps = K.random_normal(shape=shape, mean=0., stddev=1.)

        return mu + K.exp(log_var / 2) * eps