import torch


def computeKLD(mean_a_flat, logvar_a_flat, device, mean_b_flat=0.0,
               logvar_b_flat=1.0):
    """Compute the kullback-leibler divergence between two Gaussians.

        Args:
            mean_a_flat: mean of the Gaussian a.
            logvar_a_flat: Log variance of the Gaussian a.
            mean_b_flat: mean of the Gaussian b.
            logvar_b_flat: Log variance of the Gaussian b.
            device:
        Returns:
            LKD between two gausian with given parameters.
    """

    # see Appendix B from VAE paper:
    # Kingma and Welling. Auto-Encoding Variational Bayes. ICLR, 2014
    # https://arxiv.org/abs/1312.6114

    logvar_b_flat = torch.log(torch.tensor(logvar_b_flat).to(device))
    mean_b_flat = torch.tensor(mean_b_flat).to(device)

    kl = None

    if kl is None:
        var_a_flat = torch.exp(0.5 * logvar_a_flat) ** 2
        var_b_flat = torch.exp(0.5 * logvar_b_flat) ** 2

        kl = - 0.5 * torch.sum(
            1. - (var_a_flat + (mean_a_flat - mean_b_flat) ** 2) / var_b_flat
            + logvar_a_flat - logvar_b_flat
        )

    return kl