import torch


def triple_distances(positions, neighbors_j, neighbors_k):
    """
    Get all distances between atoms forming a triangle with the central atoms.
    Required e.g. for angular symmetry functions.

    Args:
        positions (torch.Tensor): Atomic positions
        neighbors_j (torch.Tensor): Indices of first neighbor in triangle
        neighbors_k (torch.Tensor): Indices of second neighbor in triangle

    Returns:
        torch.Tensor: Distance between central atom and neighbor j
        torch.Tensor: Distance between central atom and neighbor k
        torch.Tensor: Distance between neighbors

    """
    nbatch, natoms, nangles = neighbors_k.size()
    idx_m = torch.arange(nbatch, device=positions.device, dtype=torch.long)[:,
            None, None]
    # if positions.is_cuda:
    #    idx_m = idx_m.pin_memory().cuda(async=True)
    pos_j = positions[idx_m, neighbors_j[:], :]
    pos_k = positions[idx_m, neighbors_k[:], :]
    R_ij = pos_j - positions[:, :, None, :]
    R_ik = pos_k - positions[:, :, None, :]
    R_jk = pos_j - pos_k
    r_ij = torch.norm(R_ij, 2, 3)
    r_ik = torch.norm(R_ik, 2, 3)
    r_jk = torch.norm(R_jk, 2, 3)
    return r_ij, r_ik, r_jk