import torch


def moments_log(log_probs, 
        log_levels, bias_levels, 
        log_squared_levels, bias_squared_levels):
    """
    Compute mean and variance of a probability tensor.
    Note: the probability is represented in log-space. 

    Arguments:
    ----------
    log_probs: an (N+1)-th order tensor of size [dim[0], ..., dim[N-1], levels]
        Marginal distribution of the weights/bias, represented in log-space.
          Prob[w[i_0, ..., i_{N-1}] = levels[l]] 
        = exp(log_probs[i_0, ..., i_{N-1}, l]) / sum_k [exp(log_probs[i_0, ..., i_{N-1}, k])]
 
    log_levels: a vector of length [num_levels]
        Quantized values of the weights/bias, represented in log-space.
        log_levels[l] = log(levels[l] - bias_levels), 
        where bias_levels < min_k levels[k].

    bias_levels: float
        The bias to compute the log_levels.
        Typically, bias_levels = min_k levels[k] - 1.

    log_squared_levels: a vector of lenght [num_levels]
        Square of the quantized values of the weights/bias, represented in log-space.
        log_squared_levels = log(levels[l]^2 - bias_squared_levels),
        where bias_squared_levels < min_k log_squared_levels.

    bias_squared_levels: float
        The bias to compute the log_squared_levels.
        Typically, bias_squared_levels = min_k levels[k]^2 - 1. 

    Returns:
    --------
    mean: an N-th order tensor of size [dim[0], ..., dim[N-1]]
        Mean of the weights/bias.
        mean[i_0, ..., i_{N-1}] = E[w[i_0, ..., i_{N-1}]]

    var:  an N-th order tensor of size [dim[0], ..., dim[N-1]]
        Variance of the weights/bias.
         var[i_0, ..., i_{N-1}] = V[w[i_0, ..., i_{N-1}]]
    """

    # E[W] = \sum_x w(x) p(x)
    mean = torch.exp(torch.logsumexp(log_probs + log_levels, dim = -1) - 
        torch.logsumexp(log_probs, dim = -1)) + bias_levels

    # E[W^2] = \sum_x w(x)^2 p(x)
    second_moment = torch.exp(torch.logsumexp(log_probs + log_squared_levels, dim = -1) - 
        torch.logsumexp(log_probs, dim = -1)) + bias_squared_levels

    # V[W] = E[W^2] - E[W]^2
    var  = second_moment - mean ** 2

    # returns a pair of N-th order tensors
    return (mean, var)