import tensorflow as tf
from tensorflow.keras import losses
from teras.losses.gain import (generator_loss as generator_pretraining_loss,
                               discriminator_loss)


# The generator pretraining loss is the same as generator_loss
# in the GAIN model because in PC-GAIN we use a GAIN model for
# pretraining.

# In PC-GAIN, the discriminator loss is also exactly the same
# as GAIN's discriminator loss

def generator_loss(real_samples=None,
                   generated_samples=None,
                   discriminator_pred=None,
                   mask=None,
                   alpha=200,
                   beta=100,
                   classifier_pred=None):
    """
    Generator loss used during the main training phase (post-pretraining)
    of PC-GAIN architecture.
    It is similar to the generator loss used in the pretraining stage
    except for an additional `Information Entropy Loss` that is calculated
    for the Classifier's predictions and weighted by the `beta` parameter.

    Args:
        real_samples: Samples drawn from the input dataset.
        generated_samples: Samples generated by the Generator.
        discriminator_pred: Predictions made by Discriminator,
            on the real and generated samples combined.
        mask: Mask created from the raw data batch with missing values,
            it helps identify which values were original missing and
            have been imputed and filter out generated values that were
            not missing in the batch.
            Important to remember that, the Generator generates all values
            for each sample, but we only want values in place of those that
            were originally missing. So we need some way for this identification
            and filteration. This is where this mask comes in.
            In the case of loss computation, it helps apply and combined relevant
            functions and values.
        alpha: Hyperparamter that controls how much the MSE loss contributes to
            the overall loss. It is directly proportional to the MSE loss.
            Important to keep in mind that the overall loss is made up of three losses,
            `Cross Entropy loss`, `MSE loss` and `Info entropy loss`
        beta: Hyperparamter that controls how much the Info Entropy loss contributes to
            the overall loss. It is directly proportional to the Info Entropy loss.
        classifier_pred: Softmax probabilities predicted by the classifier for the
            given batch of fused together real and generated samples.
    """
    cross_entropy_loss = -tf.reduce_mean((1 - mask) * tf.math.log(discriminator_pred + 1e-8))
    mse_loss = losses.MSE(y_true=(mask * real_samples),
                                  y_pred=(mask * generated_samples))
    info_entropy_loss = -tf.reduce_mean(classifier_pred * tf.math.log(classifier_pred + 1e-8))
    loss = cross_entropy_loss + (alpha * mse_loss) + (beta * info_entropy_loss)
    return loss
