# @author lucasmiranda42
# encoding: utf-8
# module deepof

"""

deep autoencoder models for unsupervised pose detection.
- VQ-VAE: a variational autoencoder with a vector quantization latent-space (https://arxiv.org/abs/1711.00937).
- GM-VAE: a variational autoencoder with a Gaussian mixture latent-space.

"""

import tcn
import tensorflow as tf
import tensorflow_probability as tfp
from tensorflow.keras import Input, Model
from tensorflow.keras.initializers import he_uniform
from tensorflow.keras.layers import Dense, GRU, RepeatVector, TimeDistributed
from tensorflow.keras.layers import LayerNormalization, Bidirectional
from tensorflow.keras.optimizers import Nadam

import deepof.unsupervised_utils
from deepof import unsupervised_utils

tfb = tfp.bijectors
tfd = tfp.distributions
tfpl = tfp.layers

from sklearn.mixture import GaussianMixture

# noinspection PyCallingNonCallable
def get_recurrent_encoder(
    input_shape, latent_dim, gru_unroll=False, bidirectional_merge="concat"
):
    """

    Returns a deep neural network capable of encoding the motion tracking instances into a vector ready to be fed to
    one of the provided structured latent spaces.

    Args:
        input_shape (tuple): shape of the input data
        latent_dim (int): dimension of the latent space
        gru_unroll (bool): whether to unroll the GRU layers. Defaults to False.
        bidirectional_merge (str): how to merge the forward and backward GRU layers. Defaults to "concat".

    Returns:
        keras.Model: a keras model that can be trained to encode motion tracking instances into a vector.

    """

    # Define and instantiate encoder
    x = Input(shape=input_shape)
    encoder = tf.keras.layers.Conv1D(
        filters=2 * latent_dim,
        kernel_size=5,
        strides=1,  # Increased strides yield shorter sequences
        padding="same",
        activation="relu",
        kernel_initializer=he_uniform(),
        use_bias=False,
    )(x)
    encoder = tf.keras.layers.Masking(mask_value=0.0)(encoder)
    encoder = Bidirectional(
        GRU(
            2 * latent_dim,
            activation="tanh",
            recurrent_activation="sigmoid",
            return_sequences=True,
            unroll=gru_unroll,
            use_bias=True,
        ),
        merge_mode=bidirectional_merge,
    )(encoder)
    encoder = LayerNormalization()(encoder)
    encoder = Bidirectional(
        GRU(
            latent_dim,
            activation="tanh",
            recurrent_activation="sigmoid",
            return_sequences=False,
            unroll=gru_unroll,
            use_bias=True,
        ),
        merge_mode=bidirectional_merge,
    )(encoder)
    encoder = LayerNormalization()(encoder)
    encoder_output = tf.keras.layers.Dense(latent_dim, kernel_initializer="he_uniform")(
        encoder
    )

    return Model(x, encoder_output, name="recurrent_encoder")


# noinspection PyCallingNonCallable
def get_recurrent_decoder(
    input_shape, latent_dim, gru_unroll=False, bidirectional_merge="concat"
):
    """

    Returns a deep neural network capable of decoding the structured latent space generated by one of the compatible
    classes into a sequence of motion tracking instances, either reconstructing the original
    input, or generating new data from given clusters.

    Args:
        input_shape (tuple): shape of the input data
        latent_dim (int): dimensionality of the latent space
        gru_unroll (bool): whether to unroll the GRU layers. Defaults to False.
        bidirectional_merge (str): how to merge the forward and backward GRU layers. Defaults to "concat".

    Returns:
        keras.Model: a keras model that can be trained to decode the latent space into a series of motion tracking
        sequences.

    """

    # Define and instantiate generator
    g = Input(shape=latent_dim)  # Decoder input, shaped as the latent space
    x = Input(shape=input_shape)  # Encoder input, used to generate an output mask
    validity_mask = tf.math.logical_not(tf.reduce_all(x == 0.0, axis=2))

    generator = RepeatVector(input_shape[0])(g)
    generator = Bidirectional(
        GRU(
            latent_dim,
            activation="tanh",
            recurrent_activation="sigmoid",
            return_sequences=True,
            unroll=gru_unroll,
            use_bias=True,
        ),
        merge_mode=bidirectional_merge,
    )(generator, mask=validity_mask)
    generator = LayerNormalization()(generator)
    generator = Bidirectional(
        GRU(
            2 * latent_dim,
            activation="tanh",
            recurrent_activation="sigmoid",
            return_sequences=True,
            unroll=gru_unroll,
            use_bias=True,
        ),
        merge_mode=bidirectional_merge,
    )(generator)
    generator = LayerNormalization()(generator)
    generator = tf.keras.layers.Conv1D(
        filters=2 * latent_dim,
        kernel_size=5,
        strides=1,
        padding="same",
        activation="relu",
        kernel_initializer=he_uniform(),
        use_bias=False,
    )(generator)
    generator = LayerNormalization()(generator)

    x_decoded = deepof.unsupervised_utils.ProbabilisticDecoder(input_shape)(
        [generator, validity_mask]
    )

    return Model([g, x], x_decoded, name="recurrent_decoder")


def get_TCN_encoder(
    input_shape,
    latent_dim,
    conv_filters=64,
    kernel_size=4,
    conv_stacks=2,
    conv_dilations=(1, 2, 4, 8),
    padding="causal",
    use_skip_connections=True,
    dropout_rate=0,
    activation="relu",
):
    """

    Returns a Temporal Convolutional Network (TCN) encoder that can be used to encode motion tracking instances into a
    vector. Each layer contains a residual block with a convolutional layer and a skip connection. See the following
    paper for more details: https://arxiv.org/pdf/1803.01271.pdf

    Args:
        input_shape: shape of the input data
        latent_dim: dimensionality of the latent space
        conv_filters: number of filters in the TCN layers
        kernel_size: size of the convolutional kernels
        conv_stacks: number of TCN layers
        conv_dilations: list of dilation factors for each TCN layer
        padding: padding mode for the TCN layers
        use_skip_connections: whether to use skip connections between TCN layers
        dropout_rate: dropout rate for the TCN layers
        activation: activation function for the TCN layers

    Returns:
        keras.Model: a keras model that can be trained to encode a sequence of motion tracking instances into a latent
        space using temporal convolutional networks.

    """
    x = Input(input_shape)

    encoder = tcn.TCN(
        conv_filters,
        kernel_size,
        conv_stacks,
        conv_dilations,
        padding,
        use_skip_connections,
        dropout_rate,
        return_sequences=True,
        activation=activation,
        kernel_initializer="random_normal",
        use_batch_norm=True,
    )(x)

    encoder = tf.keras.layers.Flatten()(encoder)
    encoder = tf.keras.layers.Dense(2 * latent_dim, activation="relu")(encoder)
    encoder = tf.keras.layers.BatchNormalization()(encoder)
    encoder = Dense(latent_dim, activation="relu")(encoder)
    encoder = tf.keras.layers.BatchNormalization()(encoder)
    encoder = tf.keras.layers.Dense(latent_dim)(encoder)
    encoder = tf.keras.layers.BatchNormalization()(encoder)

    return Model(x, encoder, name="TCN_encoder")


def get_TCN_decoder(
    input_shape,
    latent_dim,
    conv_filters=64,
    kernel_size=4,
    conv_stacks=2,
    conv_dilations=(8, 4, 2, 1),
    padding="causal",
    use_skip_connections=True,
    dropout_rate=0,
    activation="relu",
):
    """

    Returns a Temporal Convolutional Network (TCN) decoder that can be used to decode a latent space into a sequence of
    motion tracking instances. Each layer contains a residual block with a convolutional layer and a skip connection. See
    the following paper for more details: https://arxiv.org/pdf/1803.01271.pdf

    Args:
        input_shape: shape of the input data
        latent_dim: dimensionality of the latent space
        conv_filters: number of filters in the TCN layers
        kernel_size: size of the convolutional kernels
        conv_stacks: number of TCN layers
        conv_dilations: list of dilation factors for each TCN layer
        padding: padding mode for the TCN layers
        use_skip_connections: whether to use skip connections between TCN layers
        dropout_rate: dropout rate for the TCN layers
        activation: activation function for the TCN layers

    Returns:
        keras.Model: a keras model that can be trained to decode a latent space into a sequence of motion tracking
        instances using temporal convolutional networks.

    """

    # Define and instantiate generator
    g = Input(shape=latent_dim)  # Decoder input, shaped as the latent space
    x = Input(shape=input_shape)  # Encoder input, used to generate an output mask
    validity_mask = tf.math.logical_not(tf.reduce_all(x == 0.0, axis=2))

    generator = tf.keras.layers.Dense(latent_dim)(g)
    generator = tf.keras.layers.BatchNormalization()(generator)
    generator = tf.keras.layers.Dense(2 * latent_dim, activation="relu")(generator)
    generator = tf.keras.layers.BatchNormalization()(generator)
    generator = tf.keras.layers.Dense(4 * latent_dim, activation="relu")(generator)
    generator = tf.keras.layers.BatchNormalization()(generator)
    generator = tf.keras.layers.RepeatVector(input_shape[0])(generator)

    generator = tcn.TCN(
        conv_filters,
        kernel_size,
        conv_stacks,
        conv_dilations,
        padding,
        use_skip_connections,
        dropout_rate,
        return_sequences=True,
        activation=activation,
        kernel_initializer="random_normal",
        use_batch_norm=True,
    )(generator)

    x_decoded = deepof.unsupervised_utils.ProbabilisticDecoder(input_shape)(
        [generator, validity_mask]
    )

    return Model([g, x], x_decoded, name="TCN_decoder")


# noinspection PyCallingNonCallable
def get_transformer_encoder(
    input_shape: tuple,
    latent_dim: int,
    num_layers: int = 2,
    num_heads: int = 8,
    dff: int = 128,
    dropout_rate: float = 0.1,
):
    """

    Transformer encoder. Based on https://www.tensorflow.org/text/tutorials/transformer.
    Adapted according to https://academic.oup.com/gigascience/article/8/11/giz134/5626377?login=true
    and https://arxiv.org/abs/1711.03905.

    Args:
        input_shape (tuple): shape of the input data
        latent_dim (int): dimensionality of the latent space
        num_layers (int): number of transformer layers to include
        num_heads (int): number of heads of the multi-head-attention layers used on the transformer encoder
        dff (int): dimensionality of the token embeddings
        dropout_rate (float): dropout rate

    """

    x = tf.keras.layers.Input(input_shape)

    transformer_embedding = deepof.unsupervised_utils.TransformerEncoder(
        num_layers=num_layers,
        seq_dim=input_shape[-1],
        key_dim=input_shape[-1],
        num_heads=num_heads,
        dff=dff,
        maximum_position_encoding=input_shape[0],
        rate=dropout_rate,
    )(x, training=False)

    encoder = tf.reshape(transformer_embedding, [-1, input_shape[0] * input_shape[1]])
    encoder = tf.keras.layers.Dense(2 * latent_dim, activation="relu")(encoder)
    encoder = tf.keras.layers.BatchNormalization()(encoder)
    encoder = tf.keras.layers.Dense(latent_dim, activation="relu")(encoder)
    encoder = tf.keras.layers.BatchNormalization()(encoder)
    encoder = tf.keras.layers.Dense(latent_dim)(encoder)
    encoder = tf.keras.layers.BatchNormalization()(encoder)

    return tf.keras.models.Model(x, encoder, name="transformer_encoder")


def get_transformer_decoder(
    input_shape, latent_dim, num_layers=2, num_heads=8, dff=128, dropout_rate=0.1
):
    """

    Transformer decoder. Based on https://www.tensorflow.org/text/tutorials/transformer.
    Adapted according to https://academic.oup.com/gigascience/article/8/11/giz134/5626377?login=true
    and https://arxiv.org/abs/1711.03905.

    Args:
        input_shape (tuple): shape of the input data
        latent_dim (int): dimensionality of the latent space
        num_layers (int): number of transformer layers to include
        num_heads (int): number of heads of the multi-head-attention layers used on the transformer encoder
        dff (int): dimensionality of the token embeddings
        dropout_rate (float): dropout rate

    """

    x = tf.keras.layers.Input(input_shape)
    g = tf.keras.layers.Input([latent_dim])
    validity_mask = tf.math.logical_not(tf.reduce_all(x == 0.0, axis=2))

    generator = tf.keras.layers.Dense(latent_dim)(g)
    generator = tf.keras.layers.BatchNormalization()(generator)
    generator = tf.keras.layers.Dense(2 * latent_dim, activation="relu")(generator)
    generator = tf.keras.layers.BatchNormalization()(generator)
    generator = tf.keras.layers.Dense(4 * latent_dim, activation="relu")(generator)
    generator = tf.keras.layers.BatchNormalization()(generator)
    generator = tf.keras.layers.RepeatVector(input_shape[0])(generator)

    # Get masks for generated output
    _, look_ahead_mask, padding_mask = deepof.unsupervised_utils.create_masks(generator)

    (
        transformer_embedding,
        attention_weights,
    ) = deepof.unsupervised_utils.TransformerDecoder(
        num_layers=num_layers,
        seq_dim=input_shape[-1],
        key_dim=input_shape[-1],
        num_heads=num_heads,
        dff=dff,
        maximum_position_encoding=input_shape[0],
        rate=dropout_rate,
    )(
        x,
        generator,
        training=False,
        look_ahead_mask=look_ahead_mask,
        padding_mask=padding_mask,
    )

    x_decoded = deepof.unsupervised_utils.ProbabilisticDecoder(input_shape)(
        [transformer_embedding, validity_mask]
    )

    return tf.keras.models.Model(
        [g, x], [x_decoded, attention_weights], name="transformer_decoder"
    )


class VectorQuantizer(tf.keras.models.Model):
    """

    Vector quantizer layer, which quantizes the input vectors into a fixed number of clusters using L2 norm. Based on
    https://arxiv.org/pdf/1509.03700.pdf, and adapted for clustering using https://arxiv.org/abs/1806.02199.
    Implementation based on https://keras.io/examples/generative/vq_vae/.

    """

    def __init__(
        self, n_components, embedding_dim, beta, kmeans_loss: float = 0.0, **kwargs
    ):
        """

        Initializes the VQ layer.

        Args:
            n_components (int): number of embeddings to use
            embedding_dim (int): dimensionality of the embeddings
            beta (float): beta value for the loss function
            kmeans_loss (float): regularization parameter for the Gram matrix
            **kwargs: additional arguments for the parent class

        """

        super(VectorQuantizer, self).__init__(**kwargs)
        self.embedding_dim = embedding_dim
        self.n_components = n_components
        self.beta = beta
        self.kmeans = kmeans_loss

        # Initialize the VQ codebook
        w_init = tf.random_uniform_initializer()
        self.codebook = tf.Variable(
            initial_value=w_init(
                shape=(self.embedding_dim, self.n_components), dtype="float32"
            ),
            trainable=True,
            name="vqvae_codebook",
        )

    def call(self, x):  # pragma: no cover
        """

        Computes the VQ layer.

        Args:
            x (tf.Tensor): input tensor

        Returns:
                x (tf.Tensor): output tensor

        """

        # Compute input shape and flatten, keeping the embedding dimension intact
        input_shape = tf.shape(x)

        # Add a disentangling penalty to the embeddings
        if self.kmeans:
            kmeans_loss = unsupervised_utils.compute_kmeans_loss(
                x, weight=self.kmeans, batch_size=input_shape[0]
            )
            self.add_loss(kmeans_loss)
            self.add_metric(kmeans_loss, name="kmeans_loss")

        flattened = tf.reshape(x, [-1, self.embedding_dim])

        # Quantize input using the codebook
        encoding_indices = tf.cast(
            self.get_code_indices(flattened, return_soft_counts=False), tf.int32
        )
        soft_counts = self.get_code_indices(flattened, return_soft_counts=True)

        encodings = tf.one_hot(encoding_indices, self.n_components)

        quantized = tf.matmul(encodings, self.codebook, transpose_b=True)
        quantized = tf.reshape(quantized, input_shape)

        # Compute vector quantization loss, and add it to the layer
        commitment_loss = self.beta * tf.reduce_mean(
            (tf.stop_gradient(quantized) - x) ** 2
        )
        codebook_loss = tf.reduce_mean((quantized - tf.stop_gradient(x)) ** 2)
        self.add_loss(commitment_loss + codebook_loss)

        # Straight-through estimator (copy gradients through the undiferentiable layer)
        # This approach has been reported to have issues for clustering, so we use add an extra
        # reconstruction loss to ensure that the gradients can flow through the encoder.
        # quantized = x + tf.stop_gradient(quantized - x)

        return quantized, soft_counts

    # noinspection PyTypeChecker
    def get_code_indices(
        self, flattened_inputs, return_soft_counts=False
    ):  # pragma: no cover
        """

        Getter for the code indices at any given time.

        Args:
            flattened_inputs (tf.Tensor): flattened input tensor (encoder output)
            return_soft_counts (bool): whether to return soft counts based on the distance to the codes, instead of
            the code indices

        Returns:
            encoding_indices (tf.Tensor): code indices tensor with cluster assignments.

        """
        # Compute L2-norm distance between inputs and codes at a given time
        similarity = tf.matmul(flattened_inputs, self.codebook)
        distances = (
            tf.reduce_sum(flattened_inputs ** 2, axis=1, keepdims=True)
            + tf.reduce_sum(self.codebook ** 2, axis=0)
            - 2 * similarity
        )

        if return_soft_counts:
            # Compute soft counts based on the distance to the codes
            similarity = (1 / distances) ** 2
            soft_counts = similarity / tf.expand_dims(
                tf.reduce_sum(similarity, axis=1), axis=1
            )
            return soft_counts

        # Return index of the closest code
        encoding_indices = tf.argmin(distances, axis=1)
        return encoding_indices


# noinspection PyCallingNonCallable
def get_vqvae(
    input_shape: tuple,
    latent_dim: int,
    n_components: int,
    beta: float = 1.0,
    kmeans_loss: float = 0.0,
    encoder_type: str = "recurrent",
):
    """

    Builds a Vector-Quantization variational autoencoder (VQ-VAE) model, adapted to the DeepOF setting.

    Args:
        input_shape (tuple): shape of the input to the encoder.
        latent_dim (int): dimension of the latent space.
        n_components (int): number of embeddings in the embedding layer.
        beta (float): beta parameter of the VQ loss.
        kmeans_loss (float): regularization parameter for the Gram matrix.
        encoder_type (str): type of encoder to use. Cab be set to "recurrent" (default), "TCN", or "transformer".

    Returns:
        encoder (tf.keras.Model): connected encoder of the VQ-VAE model.
        Outputs a vector of shape (latent_dim,).
        decoder (tf.keras.Model): connected decoder of the VQ-VAE model.
        quantizer (tf.keras.Model): connected embedder layer of the VQ-VAE model.
        Outputs cluster indices of shape (batch_size,).
        vqvae (tf.keras.Model): complete VQ VAE model.

    """
    vq_layer = VectorQuantizer(
        n_components,
        latent_dim,
        beta=beta,
        kmeans_loss=kmeans_loss,
        name="vector_quantizer",
    )

    if encoder_type == "recurrent":
        encoder = get_recurrent_encoder(
            input_shape=input_shape[1:], latent_dim=latent_dim
        )
        decoder = get_recurrent_decoder(
            input_shape=input_shape[1:], latent_dim=latent_dim
        )

    elif encoder_type == "TCN":
        encoder = get_TCN_encoder(input_shape=input_shape[1:], latent_dim=latent_dim)
        decoder = get_TCN_decoder(input_shape=input_shape[1:], latent_dim=latent_dim)

    elif encoder_type == "transformer":
        encoder = get_transformer_encoder(input_shape[1:], latent_dim=latent_dim)
        decoder = get_transformer_decoder(input_shape[1:], latent_dim=latent_dim)

    # Connect encoder and quantizer
    inputs = tf.keras.layers.Input(input_shape[1:], name="encoder_input")
    encoder_outputs = encoder(inputs)
    quantized_latents, soft_counts = vq_layer(encoder_outputs)

    # Connect full models
    encoder = tf.keras.Model(inputs, encoder_outputs, name="encoder")
    quantizer = tf.keras.Model(inputs, quantized_latents, name="quantizer")
    soft_quantizer = tf.keras.Model(inputs, soft_counts, name="soft_quantizer")
    vqvae = tf.keras.Model(
        quantizer.inputs, decoder([quantizer.outputs, inputs]), name="VQ-VAE"
    )

    models = [encoder, decoder, quantizer, soft_quantizer, vqvae]

    return models


class VQVAE(tf.keras.models.Model):
    """

    VQ-VAE model adapted to the DeepOF setting.

    """

    def __init__(
        self,
        input_shape: tuple,
        latent_dim: int = 16,
        n_components: int = 15,
        beta: float = 1.0,
        kmeans_loss: float = 0.0,
        encoder_type: str = "recurrent",
        **kwargs,
    ):
        """

        Initializes a VQ-VAE model.

        Args:
            input_shape (tuple): Shape of the input to the full model.
            latent_dim (int): Dimensionality of the latent space.
            n_components (int): Number of embeddings (clusters) in the embedding layer.
            beta (float): Beta parameter of the VQ loss, as described in the original VQVAE paper.
            kmeans_loss (float): Regularization parameter for the Gram matrix.
            encoder_type (str): Type of encoder to use. Can be set to "recurrent" (default), "TCN", or "transformer".
            **kwargs: Additional keyword arguments.

        """

        super(VQVAE, self).__init__(**kwargs)
        self.seq_shape = input_shape
        self.latent_dim = latent_dim
        self.n_components = n_components
        self.beta = beta
        self.kmeans = kmeans_loss
        self.encoder_type = encoder_type

        # Define VQ_VAE model
        (
            self.encoder,
            self.decoder,
            self.quantizer,
            self.soft_quantizer,
            self.vqvae,
        ) = get_vqvae(
            self.seq_shape,
            self.latent_dim,
            self.n_components,
            self.beta,
            self.kmeans,
            self.encoder_type,
        )

        # Define metrics to track
        self.total_loss_tracker = tf.keras.metrics.Mean(name="total_loss")
        self.encoding_reconstruction_loss_tracker = tf.keras.metrics.Mean(
            name="encoding_reconstruction_loss"
        )
        self.reconstruction_loss_tracker = tf.keras.metrics.Mean(
            name="reconstruction_loss"
        )
        self.vq_loss_tracker = tf.keras.metrics.Mean(name="vq_loss")
        self.cluster_population = tf.keras.metrics.Mean(
            name="number_of_populated_clusters"
        )
        self.val_total_loss_tracker = tf.keras.metrics.Mean(name="total_loss")
        self.val_encoding_reconstruction_loss_tracker = tf.keras.metrics.Mean(
            name="encoding_reconstruction_loss"
        )
        self.val_reconstruction_loss_tracker = tf.keras.metrics.Mean(
            name="reconstruction_loss"
        )
        self.val_vq_loss_tracker = tf.keras.metrics.Mean(name="vq_loss")
        self.val_cluster_population = tf.keras.metrics.Mean(
            name="number_of_populated_clusters"
        )

    @tf.function
    def call(self, inputs, **kwargs):
        return self.vqvae(inputs, **kwargs)

    @property
    def metrics(self):  # pragma: no cover
        metrics = [
            self.total_loss_tracker,
            self.encoding_reconstruction_loss_tracker,
            self.reconstruction_loss_tracker,
            self.vq_loss_tracker,
            self.cluster_population,
            self.val_total_loss_tracker,
            self.val_encoding_reconstruction_loss_tracker,
            self.val_reconstruction_loss_tracker,
            self.val_vq_loss_tracker,
            self.val_cluster_population,
        ]

        return metrics

    @tf.function
    def train_step(self, data):  # pragma: no cover
        """

        Performs a training step.

        """

        # Unpack data, repacking labels into a generator
        x, y = data
        if not isinstance(y, tuple):
            y = [y]
        y = (labels for labels in y)

        with tf.GradientTape() as tape:
            # Get outputs from the full model
            encoding_reconstructions = self.vqvae(x, training=True)
            reconstructions = self.decoder(
                [self.encoder(x, training=True), x], training=True
            )

            # Get rid of the attention scores that the transformer decoder outputs
            if self.encoder_type == "transformer":
                encoding_reconstructions = encoding_reconstructions[0]
                reconstructions = reconstructions[0]

            # Compute losses
            reconstruction_labels = next(y)
            encoding_reconstruction_loss = -tf.reduce_mean(
                encoding_reconstructions.log_prob(reconstruction_labels)
            )
            reconstruction_loss = -tf.reduce_mean(
                reconstructions.log_prob(reconstruction_labels)
            )

            total_loss = (
                encoding_reconstruction_loss
                + reconstruction_loss
                + sum(self.vqvae.losses)
            )

        # Backpropagation
        grads = tape.gradient(total_loss, self.vqvae.trainable_variables)
        self.optimizer.apply_gradients(zip(grads, self.vqvae.trainable_variables))

        # Compute populated clusters
        unique_indices = tf.unique(
            tf.reshape(tf.argmax(self.soft_quantizer(x), axis=1), [-1])
        ).y
        populated_clusters = tf.shape(unique_indices)[0]

        # Track losses
        self.total_loss_tracker.update_state(total_loss)
        self.encoding_reconstruction_loss_tracker.update_state(
            encoding_reconstruction_loss
        )
        self.reconstruction_loss_tracker.update_state(reconstruction_loss)
        self.vq_loss_tracker.update_state(sum(self.vqvae.losses))
        self.cluster_population.update_state(populated_clusters)

        # Log results (coupled with TensorBoard)
        log_dict = {
            "total_loss": self.total_loss_tracker.result(),
            "encoding_reconstruction_loss": self.encoding_reconstruction_loss_tracker.result(),
            "reconstruction_loss": self.reconstruction_loss_tracker.result(),
            "vq_loss": self.vq_loss_tracker.result(),
            "number_of_populated_clusters": self.cluster_population.result(),
        }

        return {**log_dict, **{met.name: met.result() for met in self.vqvae.metrics}}

    @tf.function
    def test_step(self, data):  # pragma: no cover
        """

        Performs a test step.

        """

        # Unpack data, repacking labels into a generator
        x, y = data
        if not isinstance(y, tuple):
            y = [y]
        y = (labels for labels in y)

        # Get outputs from the full model
        encoding_reconstructions = self.vqvae(x, training=False)
        reconstructions = self.decoder(
            [self.encoder(x, training=False), x], training=False
        )

        # Get rid of the attention scores that the transformer decoder outputs
        if self.encoder_type == "transformer":
            encoding_reconstructions = encoding_reconstructions[0]
            reconstructions = reconstructions[0]

        # Compute losses
        reconstruction_labels = next(y)
        encoding_reconstruction_loss = -tf.reduce_mean(
            encoding_reconstructions.log_prob(reconstruction_labels)
        )
        reconstruction_loss = -tf.reduce_mean(
            reconstructions.log_prob(reconstruction_labels)
        )
        total_loss = (
            encoding_reconstruction_loss + reconstruction_loss + sum(self.vqvae.losses)
        )

        # Compute populated clusters
        unique_indices = tf.unique(
            tf.reshape(tf.argmax(self.soft_quantizer(x), axis=1), [-1])
        ).y
        populated_clusters = tf.shape(unique_indices)[0]

        # Track losses
        self.val_total_loss_tracker.update_state(total_loss)
        self.val_encoding_reconstruction_loss_tracker.update_state(
            encoding_reconstruction_loss
        )
        self.val_reconstruction_loss_tracker.update_state(reconstruction_loss)
        self.val_vq_loss_tracker.update_state(sum(self.vqvae.losses))
        self.val_cluster_population.update_state(populated_clusters)

        # Log results (coupled with TensorBoard)
        log_dict = {
            "total_loss": self.val_total_loss_tracker.result(),
            "encoding_reconstruction_loss": self.val_encoding_reconstruction_loss_tracker.result(),
            "reconstruction_loss": self.val_reconstruction_loss_tracker.result(),
            "vq_loss": self.val_vq_loss_tracker.result(),
            "number_of_populated_clusters": self.val_cluster_population.result(),
        }

        return {**log_dict, **{met.name: met.result() for met in self.vqvae.metrics}}


class GaussianMixtureLatent(tf.keras.models.Model):
    """

    Gaussian Mixture probabilistic latent space model. Used to represent the embedding of motion tracking data in a
    mixture of Gaussians with a specified number of components, with means, covariances and weights specified by
    neural network layers.

    """

    def __init__(
        self,
        input_shape: tuple,
        n_components: int,
        latent_dim: int,
        batch_size: int,
        kl_warmup: int = 5,
        kl_annealing_mode: str = "linear",
        mc_kl: int = 1000,
        mmd_warmup: int = 15,
        mmd_annealing_mode: str = "linear",
        kmeans_loss: float = 0.0,
        reg_cluster_variance: bool = False,
        **kwargs,
    ):
        """

        Initializes the Gaussian Mixture Latent layer.

        Args:
            input_shape (tuple): shape of the input data
            n_components (int): number of components in the Gaussian mixture.
            latent_dim (int): dimensionality of the latent space.
            batch_size (int): batch size for training.
            kl_warmup (int): number of epochs to warm up the KL divergence.
            kl_annealing_mode (str): mode to use for annealing the KL divergence. Must be one of "linear" and "sigmoid".
            mc_kl (int): number of Monte Carlo samples to use for computing the KL divergence.
            mmd_warmup (int): number of epochs to warm up the MMD.
            mmd_annealing_mode (str): mode to use for annealing the MMD. Must be one of "linear" and "sigmoid".
            kmeans_loss (float): weight of the Gram matrix regularization loss.
            reg_cluster_variance (bool): whether to penalize uneven cluster variances in the latent space.
            **kwargs: keyword arguments passed to the parent class

        """

        super(GaussianMixtureLatent, self).__init__(**kwargs)
        self.seq_shape = input_shape
        self.n_components = n_components
        self.latent_dim = latent_dim
        self.batch_size = batch_size
        self.kl_warmup = kl_warmup
        self.kl_annealing_mode = kl_annealing_mode
        self.mc_kl = mc_kl
        self.mmd_warmup = mmd_warmup
        self.mmd_annealing_mode = mmd_annealing_mode
        self.kmeans = kmeans_loss
        self.optimizer = Nadam(learning_rate=1e-3, clipvalue=0.75)
        self.reg_cluster_variance = reg_cluster_variance

        # Initialize layers
        self.z_cat = Dense(
            self.n_components,
            name="cluster_assignment",
            kernel_initializer="random_normal",
            activation="linear",
            activity_regularizer=None,
        )
        self.z_gauss_mean = Dense(
            tfpl.IndependentNormal.params_size(self.latent_dim * self.n_components)
            // 2,
            name="cluster_means",
            activation="linear",
            activity_regularizer=None,
            kernel_initializer=he_uniform(),
        )
        self.z_gauss_var = Dense(
            tfpl.IndependentNormal.params_size(self.latent_dim * self.n_components)
            // 2,
            name="cluster_variances",
            activation="linear",
            kernel_regularizer=(
                tf.keras.regularizers.l2(0.01) if self.reg_cluster_variance else None
            ),
        )
        self.latent_distribution = tfpl.MixtureNormal(
            self.n_components,
            self.latent_dim,
            convert_to_tensor_fn="sample",
            name="encoding_distribution",
        )

        # Initialize the Gaussian Mixture prior with the specified number of components
        self.prior = tfd.MixtureSameFamily(
            mixture_distribution=tfd.categorical.Categorical(
                probs=tf.ones(self.n_components) / self.n_components
            ),
            components_distribution=tfd.Independent(
                tfd.Normal(
                    loc=tf.random.uniform([self.n_components, self.latent_dim]),
                    scale=tf.ones([self.n_components, self.latent_dim])
                    / (tf.math.sqrt(tf.cast(self.n_components, tf.float32))),
                    name="prior_scales",
                ),
                reinterpreted_batch_ndims=1,
            ),
        )

        self.cluster_control_layer = unsupervised_utils.ClusterControl(
            batch_size=self.batch_size,
            n_components=self.n_components,
            encoding_dim=self.latent_dim,
            k=self.n_components,
        )

        # Initialize metric layers
        self.kl_warm_up_iters = tf.cast(
            self.kl_warmup * (self.seq_shape // self.batch_size), tf.int64
        )
        self.kl_layer = unsupervised_utils.KLDivergenceLayer(
            distribution_b=self.prior,
            test_points_fn=lambda q: q.sample(self.mc_kl * self.n_components),
            test_points_reduce_axis=[0, 1],
            iters=self.optimizer.iterations,
            warm_up_iters=self.kl_warm_up_iters,
            annealing_mode=self.kl_annealing_mode,
        )

    def update_prior(self, embeddings: tf.Tensor):
        """

        Updates the prior based on a Gaussian Mixture Model fit to the provided embeddings.

        Args:
            embeddings (tf.Tensor): embeddings to use for updating the prior.

        """

        gmm = GaussianMixture(
            n_components=self.n_components, covariance_type="diag", reg_covar=1e-5
        )
        gmm.fit(embeddings.numpy())

        self.prior = tfd.MixtureSameFamily(
            mixture_distribution=tfd.categorical.Categorical(probs=gmm.weights_),
            components_distribution=tfd.Independent(
                tfd.Normal(
                    loc=gmm.means_,
                    scale=tf.math.sqrt(gmm.covariances_),
                    name="prior_scales",
                ),
                reinterpreted_batch_ndims=1,
            ),
        )

    def call(self, inputs):  # pragma: no cover
        """

        Computes the output of the layer.

        """

        z_cat = self.z_cat(inputs)
        z_gauss_mean = self.z_gauss_mean(inputs)
        z_gauss_var = self.z_gauss_var(inputs)
        z_gauss_var = tf.keras.layers.ActivityRegularization(l2=0.01)(z_gauss_var)

        # Prepare layer to be handled by the mixture normal layer
        z_gauss = tf.keras.layers.concatenate(
            [z_cat, z_gauss_mean, z_gauss_var], axis=-1
        )
        z_gauss = tf.keras.layers.Flatten()(z_gauss)

        z = self.latent_distribution(z_gauss)

        if self.kmeans:
            kmeans_loss = unsupervised_utils.compute_kmeans_loss(
                z, weight=self.kmeans, batch_size=self.batch_size
            )
            self.add_loss(kmeans_loss)
            self.add_metric(kmeans_loss, name="kmeans_loss")

        # Update KL weight based on the current iteration
        self.kl_layer._iters = self.optimizer.iterations

        # noinspection PyCallingNonCallable
        z = self.kl_layer(z)

        # Tracks clustering metrics
        if self.n_components > 1:
            z = self.cluster_control_layer([z, z_cat])

        return z, z_cat


# noinspection PyCallingNonCallable
def get_gmvae(
    input_shape: tuple,
    latent_dim: int,
    n_components: int,
    batch_size: int = 64,
    kl_warmup: int = 15,
    kl_annealing_mode: str = "sigmoid",
    mc_kl: int = 100,
    kmeans_loss: float = 1.0,
    reg_cluster_variance: bool = False,
    encoder_type: str = "recurrent",
):
    """

    Builds a Gaussian mixture variational autoencoder (GMVAE) model, adapted to the DeepOF setting.

    Args:
            input_shape (tuple): shape of the input data
            latent_dim (int): dimensionality of the latent space.
            n_components (int): number of components in the Gaussian mixture.
            batch_size (int): batch size for training.
            kl_warmup: Number of iterations during which to warm up the KL divergence.
            kl_annealing_mode (str): mode to use for annealing the KL divergence. Must be one of "linear" and "sigmoid".
            mc_kl (int): number of Monte Carlo samples to use for computing the KL divergence.
            kmeans_loss (float): weight of the Gram matrix loss as described in deepof.unsupervised_utils.compute_kmeans_loss.
            reg_cluster_variance (bool): whether to penalize uneven cluster variances in the latent space.
            encoder_type (str): type of encoder to use. Can be set to "recurrent" (default), "TCN", or "transformer".

    Returns:
        encoder (tf.keras.Model): connected encoder of the VQ-VAE model.
        Outputs a vector of shape (latent_dim,).
        decoder (tf.keras.Model): connected decoder of the VQ-VAE model.
        grouper (tf.keras.Model): deep clustering branch of the VQ-VAE model. Outputs a vector of shape (n_components,).
        for each training instance, corresponding to the soft counts for each cluster.
        gmvae (tf.keras.Model): complete GMVAE model

    """

    if encoder_type == "recurrent":
        encoder = get_recurrent_encoder(
            input_shape=input_shape[1:], latent_dim=latent_dim
        )
        decoder = get_recurrent_decoder(
            input_shape=input_shape[1:], latent_dim=latent_dim
        )

    elif encoder_type == "TCN":
        encoder = get_TCN_encoder(input_shape=input_shape[1:], latent_dim=latent_dim)
        decoder = get_TCN_decoder(input_shape=input_shape[1:], latent_dim=latent_dim)

    elif encoder_type == "transformer":
        encoder = get_transformer_encoder(input_shape[1:], latent_dim=latent_dim)
        decoder = get_transformer_decoder(input_shape[1:], latent_dim=latent_dim)

    latent_space = GaussianMixtureLatent(
        input_shape=input_shape[0],
        n_components=n_components,
        latent_dim=latent_dim,
        batch_size=batch_size,
        kl_warmup=kl_warmup,
        kl_annealing_mode=kl_annealing_mode,
        mc_kl=mc_kl,
        kmeans_loss=kmeans_loss,
        reg_cluster_variance=reg_cluster_variance,
        name="gaussian_mixture_latent",
    )

    # Connect encoder and latent space
    inputs = Input(input_shape[1:])
    encoder_outputs = encoder(inputs)
    latent, categorical = latent_space(encoder_outputs)
    embedding = tf.keras.Model(inputs, latent, name="encoder")
    grouper = tf.keras.Model(inputs, categorical, name="grouper")

    # Connect decoder
    gmvae_outputs = decoder([embedding.outputs, inputs])

    # Instantiate fully connected model
    gmvae = tf.keras.Model(embedding.inputs, gmvae_outputs, name="GMVAE")

    return embedding, decoder, grouper, gmvae


# noinspection PyDefaultArgument,PyCallingNonCallable
class GMVAE(tf.keras.models.Model):
    """

    Gaussian Mixture Variational Autoencoder for pose motif elucidation.

    """

    def __init__(
        self,
        input_shape: tuple,
        latent_dim: int = 16,
        n_components: int = 15,
        batch_size: int = 64,
        kl_annealing_mode: str = "linear",
        kl_warmup_epochs: int = 15,
        montecarlo_kl: int = 1000,
        kmeans_loss: float = 1.0,
        reg_cat_clusters: float = 1.0,
        reg_cluster_variance: bool = False,
        encoder_type: str = "recurrent",
        **kwargs,
    ):
        """

        Initalizes a GMVAE model.

        Args:
            input_shape (tuple): Shape of the input to the full model.
            batch_size (int): Batch size for training.
            latent_dim (int): Dimensionality of the latent space.
            kl_annealing_mode (str): Annealing mode for KL annealing. Can be one of 'linear' and 'sigmoid'.
            kl_warmup_epochs (int): Number of epochs to warmup KL annealing.
            montecarlo_kl (int): Number of Monte Carlo samples for KL divergence.
            n_components (int): Number of mixture components in the latent space.
            kmeans_loss (float): weight of the gram matrix regularization loss.
            reg_cat_clusters (bool): whether to use the penalize uneven cluster membership in the latent space, by
            minimizing the KL divergence between cluster membership and a uniform categorical distribution.
            reg_cluster_variance (bool): whether to penalize uneven cluster variances in the latent space.
            encoder_type (str): type of encoder to use. Cab be set to "recurrent" (default), "TCN", or "transformer".
            **kwargs: Additional keyword arguments.

        """
        super(GMVAE, self).__init__(**kwargs)
        self.seq_shape = input_shape
        self.batch_size = batch_size
        self.latent_dim = latent_dim
        self.kl_annealing_mode = kl_annealing_mode
        self.kl_warmup = kl_warmup_epochs
        self.mc_kl = montecarlo_kl
        self.n_components = n_components
        self.optimizer = Nadam(learning_rate=1e-3, clipvalue=0.75)
        self.kmeans = kmeans_loss
        self.reg_cat_clusters = reg_cat_clusters
        self.reg_cluster_variance = reg_cluster_variance
        self.encoder_type = encoder_type

        # Define GMVAE model
        self.encoder, self.decoder, self.grouper, self.gmvae = get_gmvae(
            input_shape=self.seq_shape,
            n_components=self.n_components,
            latent_dim=self.latent_dim,
            batch_size=self.batch_size,
            kl_warmup=self.kl_warmup,
            kl_annealing_mode=self.kl_annealing_mode,
            mc_kl=self.mc_kl,
            kmeans_loss=self.kmeans,
            reg_cluster_variance=self.reg_cluster_variance,
            encoder_type=self.encoder_type,
        )
        # Propagate the optimizer to all relevant sub-models, to enable metric annealing
        self.gmvae.optimizer = self.optimizer
        self.gmvae.get_layer("gaussian_mixture_latent").optimizer = self.optimizer

        # Define metrics to track

        # Track all loss function components
        self.total_loss_tracker = tf.keras.metrics.Mean(name="total_loss")
        self.val_total_loss_tracker = tf.keras.metrics.Mean(name="total_loss")

        self.reconstruction_loss_tracker = tf.keras.metrics.Mean(
            name="reconstruction_loss"
        )
        self.val_reconstruction_loss_tracker = tf.keras.metrics.Mean(
            name="reconstruction_loss"
        )

        if self.reg_cat_clusters:
            self.cat_cluster_loss_tracker = tf.keras.metrics.Mean(
                name="cat_cluster_loss"
            )
            self.val_cat_cluster_loss_tracker = tf.keras.metrics.Mean(
                name="cat_cluster_loss"
            )

    @property
    def metrics(self):  # pragma: no cover
        metrics = [
            self.total_loss_tracker,
            self.val_total_loss_tracker,
            self.reconstruction_loss_tracker,
            self.val_reconstruction_loss_tracker,
        ]
        metrics += self.gmvae.metrics

        if self.reg_cat_clusters:
            metrics += [
                self.cat_cluster_loss_tracker,
                self.val_cat_cluster_loss_tracker,
            ]

        return metrics

    @property
    def prior(self):
        """

        Property to retrieve the current model prior

        """

        return self.gmvae.get_layer("gaussian_mixture_latent").prior

    def update_prior(self, embeddings: tf.Tensor):
        """

        Updates the current prior by fitting a Gaussian Mixture Model to the current embeddings.

        """

        return self.gmvae.get_layer("gaussian_mixture_latent").update_prior(embeddings)

    @tf.function
    def call(self, inputs, **kwargs):
        return self.gmvae(inputs, **kwargs)

    def train_step(self, data):  # pragma: no cover
        """

        Performs a training step.

        """

        # Unpack data, repacking labels into a generator
        x, y = data
        if not isinstance(y, tuple):
            y = [y]
        y = (labels for labels in y)

        with tf.GradientTape() as tape:

            # Get outputs from the full model
            outputs = self.gmvae(x, training=True)

            # Get rid of the attention scores that the transformer decoder outputs
            if self.encoder_type == "transformer":
                outputs = outputs[0]

            if isinstance(outputs, list):
                reconstructions = outputs[0]
            else:
                reconstructions = outputs

            # Update the prior to reflect the cluster distribution in the current embeddings
            embeddings = self.encoder(x, training=True)
            self.update_prior(embeddings.sample())

            # Compute losses
            seq_inputs = next(y)
            total_loss = sum(self.gmvae.losses)

            # Add a regularization term to the soft_counts, to prevent the embedding layer from
            # collapsing into a few clusters.
            if self.reg_cat_clusters:

                soft_counts = self.grouper(x, training=True)
                soft_counts_regulrization = (
                    self.reg_cat_clusters
                    * deepof.unsupervised_utils.cluster_frequencies_regularizer(
                        soft_counts=soft_counts, k=self.n_components
                    )
                )
                total_loss += soft_counts_regulrization

            # Compute reconstruction loss
            reconstruction_loss = -tf.reduce_mean(reconstructions.log_prob(seq_inputs))

            total_loss += reconstruction_loss

        # Backpropagation
        grads = tape.gradient(total_loss, self.gmvae.trainable_variables)
        self.optimizer.apply_gradients(zip(grads, self.gmvae.trainable_variables))

        # Track losses
        self.total_loss_tracker.update_state(total_loss)
        self.reconstruction_loss_tracker.update_state(reconstruction_loss)

        # Log results (coupled with TensorBoard)
        log_dict = {
            "total_loss": self.total_loss_tracker.result(),
            "reconstruction_loss": self.reconstruction_loss_tracker.result(),
        }

        if self.reg_cat_clusters:
            self.cat_cluster_loss_tracker.update_state(soft_counts_regulrization)
            log_dict["cat_cluster_loss"] = self.cat_cluster_loss_tracker.result()

        # Log to TensorBoard, both explitly and implicitly (within model) tracked metrics
        return {**log_dict, **{met.name: met.result() for met in self.gmvae.metrics}}

    # noinspection PyUnboundLocalVariable
    @tf.function
    def test_step(self, data):  # pragma: no cover
        """

        Performs a test step.

        """

        # Unpack data, repacking labels into a generator
        x, y = data
        if not isinstance(y, tuple):
            y = [y]
        y = (labels for labels in y)

        # Get outputs from the full model
        outputs = self.gmvae(x, training=False)

        # Get rid of the attention scores that the transformer decoder outputs
        if self.encoder_type == "transformer":
            outputs = outputs[0]

        if isinstance(outputs, list):
            reconstructions = outputs[0]
        else:
            reconstructions = outputs

        # Compute losses
        seq_inputs = next(y)
        total_loss = sum(self.gmvae.losses)

        # Add a regularization term to the soft_counts, to prevent the embedding layer from
        # collapsing into a few clusters.
        if self.reg_cat_clusters:
            soft_counts = self.grouper(x, training=False)
            soft_counts_regulrization = (
                self.reg_cat_clusters
                * deepof.unsupervised_utils.cluster_frequencies_regularizer(
                    soft_counts=soft_counts, k=self.n_components
                )
            )
            total_loss += soft_counts_regulrization

        # Compute reconstruction loss
        reconstruction_loss = -tf.reduce_mean(reconstructions.log_prob(seq_inputs))
        total_loss += reconstruction_loss

        # Track losses
        self.val_total_loss_tracker.update_state(total_loss)
        self.val_reconstruction_loss_tracker.update_state(reconstruction_loss)

        # Log results (coupled with TensorBoard)
        log_dict = {
            "total_loss": self.val_total_loss_tracker.result(),
            "reconstruction_loss": self.val_reconstruction_loss_tracker.result(),
        }

        if self.reg_cat_clusters:
            self.val_cat_cluster_loss_tracker.update_state(soft_counts_regulrization)
            log_dict["cat_cluster_loss"] = self.val_cat_cluster_loss_tracker.result()

        return {**log_dict, **{met.name: met.result() for met in self.gmvae.metrics}}


# noinspection PyDefaultArgument,PyCallingNonCallable
class Contrastive(tf.keras.models.Model):
    """

    Gaussian Mixture Variational Autoencoder for pose motif elucidation.

    """

    def __init__(
        self,
        input_shape: tuple,
        latent_dim: int = 16,
        n_components: int = 15,
        kmeans_loss: float = 1.0,
        encoder_type: str = "recurrent",
        **kwargs,
    ):
        """

        Initalizes a self-supervised Contrastive embedding model.

        Args:
            input_shape (tuple): Shape of the input to the full model.
            latent_dim (int): Dimensionality of the latent space.
            n_components (int): Number of mixture components in the latent space.
            kmeans_loss (float): weight of the gram matrix regularization loss.
            encoder_type (str): type of encoder to use. Cab be set to "recurrent" (default), "TCN", or "transformer".
            **kwargs: Additional keyword arguments.

        """
        super(Contrastive, self).__init__(**kwargs)
        self.seq_shape = input_shape
        self.latent_dim = latent_dim
        self.n_components = n_components
        self.optimizer = Nadam(learning_rate=1e-3, clipvalue=0.75)
        self.kmeans = kmeans_loss
        self.encoder_type = encoder_type

        # Define GMVAE model
        if encoder_type == "recurrent":
            self.encoder = get_recurrent_encoder(
                input_shape=input_shape[1:], latent_dim=latent_dim
            )

        elif encoder_type == "TCN":
            self.encoder = get_TCN_encoder(
                input_shape=input_shape[1:], latent_dim=latent_dim
            )

        elif encoder_type == "transformer":
            self.encoder = get_transformer_encoder(
                input_shape[1:], latent_dim=latent_dim
            )

        # Define metrics to track

        # Track all loss function components
        self.total_loss_tracker = tf.keras.metrics.Mean(name="total_loss")
        self.val_total_loss_tracker = tf.keras.metrics.Mean(name="total_loss")

    @property
    def metrics(self):  # pragma: no cover
        metrics = [self.total_loss_tracker, self.val_total_loss_tracker]

        return metrics

    @tf.function
    def call(self, inputs, **kwargs):
        return self.encoder(inputs, **kwargs)

    def train_step(self, x):  # pragma: no cover
        """

        Performs a training step.

        """

        with tf.GradientTape() as tape:

            # Get positive and negative pairs
            def ts_samples(mbatch, win):

                x = mbatch[:, :win]
                y = mbatch[:, -win:]

                return x, y

            pos, neg = ts_samples(x, self.window_length)

            # Compute contrastive loss
            enc_pos = self.encoder(pos, training=True)
            enc_neg = self.encoder(neg, training=True)

            # normalize projection feature vectors
            enc_pos = tf.math.l2_normalize(enc_pos, axis=1)
            enc_neg = tf.math.l2_normalize(enc_neg, axis=1)

            # loss, mean_sim = ls.dcl_loss_fn(zis, zjs, temperature, lfn)
            contrastive_loss, mean_sim, neg_sim = unsupervised_utils.select_contrastive_loss(
                enc_pos,
                enc_neg,
                similarity=self.similarity_function,
                loss_fn=self.loss_function,
                temperature=self.temperature,
                tau=self.tau,
                beta=self.beta,
                elimination_th=0,
                attraction=False,
            )

            total_loss = contrastive_loss

            # Add kmeans loss if required
            if self.kmeans:
                kmeans_loss = unsupervised_utils.compute_kmeans_loss(
                    x, weight=self.kmeans, batch_size=self.input_shape[0]
                )
                self.add_loss(kmeans_loss)
                self.add_metric(kmeans_loss, name="kmeans_loss")

                total_loss += kmeans_loss

        # Backpropagation
        grads = tape.gradient(total_loss, self.encoder.trainable_variables)
        self.optimizer.apply_gradients(zip(grads, self.encoder.trainable_variables))

        # Track losses
        self.total_loss_tracker.update_state(total_loss)

        # Log results (coupled with TensorBoard)
        log_dict = {"total_loss": self.total_loss_tracker.result()}

        # Log to TensorBoard, both explitly and implicitly (within model) tracked metrics
        return {**log_dict, **{met.name: met.result() for met in self.gmvae.metrics}}

    # noinspection PyUnboundLocalVariable
    @tf.function
    def test_step(self, x):  # pragma: no cover
        """

        Performs a test step.

        """

        # Get outputs from the full model
        outputs = self.encoder(x, training=False)

        # Get rid of the attention scores that the transformer decoder outputs
        if self.encoder_type == "transformer":
            outputs = outputs[0]

        # Compute losses
        n_pairs_loss = 0  # TODO
        total_loss = n_pairs_loss

        if self.kmeans:
            kmeans_loss = unsupervised_utils.compute_kmeans_loss(
                x, weight=self.kmeans, batch_size=self.input_shape[0]
            )
            self.add_loss(kmeans_loss)
            self.add_metric(kmeans_loss, name="kmeans_loss")

            total_loss = kmeans_loss

        # Track losses
        self.val_total_loss_tracker.update_state(total_loss)

        # Log results (coupled with TensorBoard)
        log_dict = {"total_loss": self.val_total_loss_tracker.result()}

        return {**log_dict, **{met.name: met.result() for met in self.gmvae.metrics}}
