"""
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 lamareg.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
        if isinstance(last_tensor, list):
            last_tensor = last_tensor[0]
            input_shape = list(last_tensor.shape[1:])
        else:
            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.0, stddev=1.0)
        return mu + K.exp(log_var / 2) * eps
