import torch


def nll_lorentzian_var(preds, target, gamma):
    """
        Isotropic lorentzian loss function

        :param preds: prediction values from NN of size [batch, particles, timesteps, (x,y,v_x,v_y)]
        :param target: target data of size [batch, particles, timesteps, (x,y,v_x,v_y)]
        :param gamma: The tensor for the FWHM of the distribution of size [batch, particles, timesteps, (x,y,v_x,v_y)]
        :return: variance of the loss function normalised by (batch * number of atoms)
        """
    gammasquared = gamma ** 2
    neg_log_p = torch.log(1+((preds - target) ** 2 / (gammasquared)))
    neg_log_p += torch.log(gamma)
    return (neg_log_p.sum(dim=1)/target.size(1)).var()