import torch
from typing import Optional


def dkl_diag_gaussian(
    mu_1: torch.Tensor,
    var_1: torch.Tensor,
    mu_2: Optional[torch.Tensor] = None,
    var_2: Optional[torch.Tensor] = None,
    log_flag: bool = False,
    avg: bool = True,
) -> torch.Tensor:
    """Computes the analytical KL divergence between two diagonal-covariance Gaussians.

    Consider two Gaussian distributions D_1 = N(mu_1, var_1) and D_2 = N(mu_2, var_2).
    This function will compute D_KL(D_1 || D_2). If the parameters of D_2 are none,
    then D_2 is assumed to be the standard normal distribution.

    Parameters
    ----------
    mu_1 : torch.Tensor, shape=(B, X)
        Mean of D_1.
    var_1 : torch.Tensor, shape=(B, X)
        Diagonal entries of covariance of D_1.
    mu_2 : torch.Tensor, shape=(B, X), default=None
        Mean of D_2. Optional.
    var_2 : torch.Tensor, shape=(B, X), default=None
        Diagonal entries of covariance of D_2. Optional.
    log_flag : bool, default=False
        Flag indicating whether the variances are passed as log-variances
    avg : bool, default=True
        Flag indicating whether to batch average the KL divergence.

    Returns
    -------
    dkl : torch.Tensor, shape=(1) or shape=(B, 1)
        The KL divergence.
    """
    B, X = mu_1.shape
    if log_flag:
        if mu_2 is None or var_2 is None:
            inner = -var_1 + torch.exp(var_1) + mu_1 ** 2
        else:
            mu_diff = mu_2 - mu_1
            inner = var_2 - var_1 + (torch.exp(var_1) + mu_diff ** 2) / torch.exp(var_2)
    else:
        if mu_2 is None or var_2 is None:
            inner = -torch.log(var_1) + var_1 + mu_1 ** 2

        else:
            inner = torch.log(var_2 / var_1) + (var_1 + (mu_2 - mu_1) ** 2) / var_2

    dkl = 0.5 * (torch.sum(inner, dim=-1) - X)
    if avg:
        dkl = torch.sum(dkl) / B
    return dkl