import torch
import math


def lb_mvn_mixmvn_cross_entropy(mean, var, mixmeans, mixvars, mixweights, logdelta=-20):
    """
    Lower bound of
    -E_{\mathcal{N}(mean,var)}[\sum_i w_i \mathcal{N}(mean_i,var_i)],
    using Jensen's inequality

    Parameters
    ----------
    mean : torch.Tensor
        Mean vector of diagonal Gaussian distribution
    var : torch.Tensor
        Variance vector of diagonal Gaussian distribution
    mixmeans : torch.Tensor
        Mean matrix of mixtures of
        diagonal normal distribution, of shape (nmixtures,dim)
    mixvars : torch.Tensor
        Variance matrix of mixtures of
        diagonal normal distribution, of shape (nmixtures,dim)
    mixweights : torch.Tensor
        Weights vector of mixture components
    logdelta : float
        Logarithm of regularizer term for cross entropy

    Returns
    -------
    torch.Tensor
        Lower bound of cross-entropy
    """
    #mean : (n,)
    #var : (n,)
    #mixmeans : (m,n)
    #mixvars : (m,n)
    #mixweights : (m,)
    # -\log(\sum_j (\prod_k \sqrt(2 \pi (\sigma_k^2 + \sigma_j,k^2))
    w = mixweights*torch.prod(torch.sqrt(2*math.pi*(var + mixvars)), dim=-1)
    logz = -0.5*torch.sum(((mean-mixmeans)/(var + mixvars))**2, dim=-1)
    res = -torch.log(torch.sum(w*torch.exp(logz)) + math.exp(logdelta))
    return res