import torch


def split_nbrs(nbrs,
               mol_size,
               confs_in_batch,
               conf_idx):
    """
    Get the indices of the neighbor list that correspond to conformers 
    we're keeping.
    Args:
        nbrs (torch.LongTensor): neighbor list
        mol_size (int): Number of atoms in each conformer
        confs_in_batch (int): Total number of conformers in the batch
        conf_idx (list[int]): Indices of the conformers we're keeping 
    Returns:
        tens_idx (torch.LongTensor): nbr indices of conformers we're
            keeping.
    """

    split_idx = []
    # The cutoff of the atom indices for each conformer
    cutoffs = [i * mol_size - 1 for i in range(1, confs_in_batch + 1)]

    for i in conf_idx:

        # start index of the conformer
        start = cutoffs[i] - mol_size + 1
        # end index of the conformer
        end = cutoffs[i]

        # mask so that all elements of the neighbor list are
        # >= start and <= end
        mask = (nbrs[:, 0] <= end) * (nbrs[:, 1] <= end)
        mask *= (nbrs[:, 0] >= start) * (nbrs[:, 1] >= start)

        idx = mask.nonzero().reshape(-1)
        split_idx.append(idx)

    tens_idx = torch.cat(split_idx)

    return tens_idx