import tensorflow as tf
from tensorflow.keras import losses


def discriminator_loss(discriminator_pred, mask):
    """
    Discriminator loss based on the architecture and implementation
    of GAIN as proposed and implemented by
    Jinsung Yoon et al. in the paper
    GAIN: Missing Data Imputation using Generative Adversarial Nets.

    It computes the cross entropy loss of discriminator's predictions.

    Reference(s):
        https://arxiv.org/abs/1806.02920
        https://github.com/jsyoon0823/GAIN

    Args:
        discriminator_pred: Discriminator's predictions
        mask: Mask generated from inputs.

    Returns:
        Cross entropy loss for discriminator's predictions.
    """
    loss = -tf.reduce_mean(mask * tf.math.log(discriminator_pred + 1e-8)
                           + (1 - mask) * tf.math.log(1. - discriminator_pred + 1e-8))
    return loss


def generator_loss(real_samples=None,
                   generated_samples=None,
                   discriminator_pred=None,
                   mask=None,
                   alpha=None):
    """
    Generator loss based on the architecture and implementation
    of GAIN as proposed and implemented by
    Jinsung Yoon et al. in the paper
    GAIN: Missing Data Imputation using Generative Adversarial Nets.

    This loss is made up of two components,
    cross_entropy_loss: That is computed for discriminator's predictions.
    mse_loss: That is computed between the masked original and generated samples
    These losses are combined with `alpha` determining the weight of mse loss
    Precisely, `loss` = `cross_entropy_loss` + `alpha` * `mse_loss`

    Reference(s):
        https://arxiv.org/abs/1806.02920
        https://github.com/jsyoon0823/GAIN

    Args:
        real_samples: Real samples drawn from the dataset
        generated_samples: Samples generated by Generator
        discriminator_pred: Predictions output by Discriminator
        mask: Mask generated from the inputs
        alpha: The hyperparameter alpha used to control the
            weightage during addition of two losses.

    Returns:
        Generator loss made up cross entropy and MSE loss.
    """
    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))
    loss = cross_entropy_loss + alpha * mse_loss
    return loss
