import torch


def gaussian_filter(distances, offsets, widths, centered=False):
    """
    Filtered interatomic distance values using Gaussian functions.

    B   :  Batch size
    At  :  Total number of atoms in the batch
    Nbr :  Total number of neighbors of each atom
    G   :  Filtered features

    Parameters
    ----------
    distances : torch.Tensor
        interatomic distances of (B x At x Nbr) shape.
    offsets : torch.Tensor
        offsets values of Gaussian functions.
    widths : torch.Tensor
        width values of Gaussian functions.
    centered : bool, default=False
        If True, Gaussians are centered at the origin and the offsets are used
        to as their widths.

    Returns
    -------
    filtered_distances : torch.Tensor
        filtered distances of (B x At x Nbr x G) shape.

    References
    ----------
    .. [1] https://github.com/atomistic-machine-learning/schnetpack/blob/67226795af55719a7e4565ed773881841a94d130/src/schnetpack/nn/acsf.py
    """
    if centered:
        # if Gaussian functions are centered, use offsets to compute widths
        eta = 0.5 / torch.pow(offsets, 2)
        # if Gaussian functions are centered, no offset is subtracted
        diff = distances[:, :, :, None]

    else:
        # compute width of Gaussian functions (using an overlap of 1 STDDEV)
        eta = 0.5 / torch.pow(widths, 2)
        # Use advanced indexing to compute the individual components
        diff = distances[:, :, :, None] - offsets[None, None, None, :]

    # compute smear distance values
    filtered_distances = torch.exp(-eta * torch.pow(diff, 2))
    return filtered_distances