import torch


def nll_gauss(target, predictions:list, total = True):
    """
    Calculates gaussian negative log likelihood score

    Parameters
    ----------
    target : torch.Tensor
        The true values of the target variable
    predictions : list
        - predictions[0] = expected_value, a torch.Tensor containing predicted expected values
        of the target variable
        - predictions[1] = log_variance, approx. equal to log(2*pi*sigma^2)
    total : bool, default = True
        - When total is set to True, return overall gaussian negative log likelihood loss
        - When total is set to False, return gaussian negative log likelihood loss along the horizon

    Returns
    -------
    torch.Tensor
        The gaussian negative log likelihood loss, which depending on the value of 'total'
        is either a scalar (overall loss) or 1d-array over the horizon, in which case it is
        expected to increase as we move along the horizon. Generally, lower is better.
    """

    assert len(predictions) == 2
    expected_value = predictions[0]
    log_variance = predictions[1]

    # y, y_pred, var_pred must have the same shape
    assert target.shape == expected_value.shape # target.shape = torch.Size([batchsize, horizon, # of target variables]) e.g.[64,40,1]
    assert target.shape == log_variance.shape

    squared_errors = (target - expected_value) ** 2
    if total:
        return torch.mean(squared_errors / (2 * log_variance.exp()) + 0.5 * log_variance)
    else:
        return torch.mean(squared_errors / (2 * log_variance.exp()) + 0.5 * log_variance, dim=0)