import torch


def entropy_(probs, reduction = "sum"):
    """
    Compute the entropy of a probability tensor.
    Note: the probability is represented in original space.

    Arguments:
    ----------
    probs: an (N+1)-th order tensor of size [dim[0], ..., dim[N-1], num_levels]
        Marginal distribution of the weights/bias.
        Prob[w[i_0, ..., i_{N-1}] = levels[l]] = probs[i_0, ..., i_{N-1}, l]

    reduction: str (options: "mean", "sum" or "none")
        Reduction mode on the output.
        default: "sum"

    Returns:
    --------
    entropy: a scalar / an N-th order tensor (depending on the reduction mode)
        Entropy of the probability tensor.
            "none": an N-th order tensor of [dim[0], ..., dim[N-1]].
            "sum"/"mean": a float scalar.
        Note: H(p) = - sum_w p(w) log(p(w)) (in probability space)
                   = - sum_w softmax(q(w)) log_softmax(q(w)) (in log-probability space)
    """

    # compute the elementwise entropy 
    outputs = - torch.log(probs) * probs 

    # reduce the outputs according to the reduction mode
    if reduction == "sum":
        outputs = torch.sum(outputs)
    elif reduction == "mean":
        outputs = torch.mean(outputs)
    elif reduction == "none":
        outputs = torch.sum(outputs, dim = -1)
    else:
        raise NotImplementedError

    return outputs