import torch


def batch_angle_idx(nbrs):
    """
    Given a neighbor list, find the sets of indices in the neighbor list
    corresponding to the kj and ji indices. Usually you can only do this 
    for one conformer without running out of memory -- to do it for multiple 
    conformers, use `full_angle_idx` below.
    Args:
        nbrs (torch.LongTensor): neighbor list
    Returns:
        ji_idx (torch.LongTensor): a set of indices for the neighbor list
        kj_idx (torch.LongTensor): a set of indices for the neighbor list
            such that nbrs[kj_idx[n]][0] == nbrs[ji_idx[n]][1] for any
            value of n.
    """

    all_idx = torch.stack([torch.arange(len(nbrs))] * len(nbrs)).long()
    mask = ((nbrs[:, 1] == nbrs[:, 0, None])
            * (nbrs[:, 0] != nbrs[:, 1, None]))
    ji_idx = all_idx[mask]
    kj_idx = mask.nonzero(as_tuple=False)[:, 0]

    return ji_idx, kj_idx