import torch


def mollifier_cutoff(distances, cutoff=5.0):
    """
    Infinitely differentiable cutoff based on a mollifier function.
    (See https://en.wikipedia.org/wiki/Mollifier)

    Args:
        distances (torch.Tensor): Interatomic distances (Nbatch x Nat x Nneigh)
        cutoff (float): Cutoff value, all values beyond are set to 0

    Returns:
        torch.Tensor: Tensor holding values of the cutoff function (Nbatch x Nat x Nneigh)
    """
    mask = (distances <= cutoff).float()
    exponent = 1.0 - 1.0 / (1.0 - torch.pow(distances * mask / cutoff, 2))
    cutoffs = torch.exp(exponent)
    cutoffs = cutoffs * mask
    return cutoffs