import tensorflow as tf
from tensorflow.keras import losses


def elbo_loss_tvae(data_dim,
                   real_samples=None,
                   generated_samples=None,
                   meta_data=None,
                   sigmas=None,
                   mean=None,
                   log_var=None,
                   loss_factor=2):
    """
    Evidence Lower Bound (ELBO) Loss [1] adapted for
    TVAE architecture [2] proposed by
    Lei Xu et al. in the paper,
    "Modeling Tabular data using Conditional GAN".

    Reference(s):
        [1]: https://arxiv.org/abs/1312.6114
        [2]: https://arxiv.org/abs/1907.00503

    Args:
        data_dim: Dimensionality of the input dataset.
        real_samples: Samples drawn from the input dataset.
        generated_samples: Samples generated by the decoder.
        meta_data: A namedtuple containing meta data for all
            features in the input data. This meta data is computed
            during the data transformation step and can be accessed
            from `.get_meta_data()` method of DataTransformer instance.
        sigmas: Sigma values returned by decoder
        mean: Mean values returned by encoder
        log_var: Log var values returned by encoder
        loss_factor: Loss factor that controls how much the cross entropy
            loss contributes to the overall loss.
            It is directly proportional to the cross entropy loss.

    Returns:
        Elbo loss adapted for the TVAE model.
    """
    loss = []
    cross_entropy = losses.SparseCategoricalCrossentropy(from_logits=True,
                                                         reduction=losses.Reduction.SUM)
    cont_i = 0  # numerical index
    cat_i = 0  # categorical index
    for i, relative_index in enumerate(meta_data.relative_indices_all):
        # the first k features are numerical
        # TODO: replace it with len(meta_data.numerical_features) but it will require changing the DataTransformer
        if i < len(meta_data.numerical.num_valid_clusters_all):
            # each numerical feature is of the form
            # [alpha, beta1, beta2...beta(n)] where n is the number of clusters

            # calculate alpha loss
            std = sigmas[relative_index]
            eq = real_samples[:, relative_index] - tf.nn.tanh(generated_samples[:, relative_index])
            loss_temp = tf.reduce_sum((eq ** 2 / 2 / (std ** 2)))
            loss.append(loss_temp)
            loss_temp = tf.math.log(std) * tf.cast(tf.shape(real_samples)[0], dtype=tf.float32)
            loss.append(loss_temp)

            # calculate betas loss
            num_clusters = meta_data.numerical.num_valid_clusters_all[cont_i]
            logits = generated_samples[:, relative_index + 1: relative_index + 1 + num_clusters]
            labels = tf.argmax(real_samples[:, relative_index + 1: relative_index + 1 + num_clusters], axis=-1)
            cross_entropy_loss = cross_entropy(y_pred=logits, y_true=labels)
            loss.append(cross_entropy_loss)
            cont_i += 1
        else:
            num_categories = meta_data.categorical.num_categories_all[cat_i]
            logits = generated_samples[:, relative_index: relative_index + num_categories]
            labels = tf.argmax(real_samples[:, relative_index: relative_index + num_categories], axis=-1)
            cross_entropy_loss = cross_entropy(y_pred=logits, y_true=labels)
            loss.append(cross_entropy_loss)
            cat_i += 1
    KLD = -0.5 * tf.reduce_sum(1 + log_var - mean ** 2 - tf.exp(log_var))
    loss_1 = tf.reduce_sum(loss) * loss_factor / data_dim
    loss_2 = KLD / data_dim
    final_loss = loss_1 + loss_2
    return final_loss
