import torch


def compute_angles(xyz, particle_index):
    """ Compute angles between sets of three particles.

    Parameters:
    -----------
    xyz: torch.Tensor 
        input tensor of shape :math:`(\text{frames} , \text{particles} , 3)`
    particle_index: torch.LongTensor
        particle index tensor of shape :math:`(\text{*} , 3)`

    """

    xyz_i = torch.index_select(xyz, 1, particle_index[:, 0])
    xyz_j = torch.index_select(xyz, 1, particle_index[:, 1])
    xyz_k = torch.index_select(xyz, 1, particle_index[:, 2])    

    v = xyz_i - xyz_j
    w = xyz_k - xyz_j

    v = v / torch.sqrt(torch.sum(v**2, -1, keepdim = True))
    w = w / torch.sqrt(torch.sum(w**2, -1, keepdim = True))

    inner_product = torch.sum(v*w, -1)
    angles = torch.acos(inner_product)

    return angles