import torch


def compute_dihedrals(xyz, particle_index):
    """ Compute dihedral angles between sets of four 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{*} , 4)`

    """

    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])
    xyz_l = torch.index_select(xyz, 1, particle_index[:, 3])    

    b1 = xyz_i - xyz_j
    b2 = xyz_j - xyz_k
    b3 = xyz_l - xyz_k

    b1_cross_b2 = torch.cross(b1, b2)
    b3_cross_b2 = torch.cross(b3, b2)

    cos_d = torch.norm(b2, dim = -1)*torch.sum(b1_cross_b2*b3_cross_b2, -1)
    sin_d = torch.sum(b2*torch.cross(b3_cross_b2, b1_cross_b2), -1)
    
    dihedrals = torch.atan2(sin_d, cos_d)

    return dihedrals