import torch


def _grad_spherical_harmonics_l1(xyz, m):
    """Compute the nabla of 1-1 Spherical Harmonics
    Args:
        xyz : array (Nbatch,Nelec,Nrbf,Ndim) x,y,z, of (Point - Center)
        m : second quantum number (-1,0,1)
    Returns
        \nabla Y0-1 = \sqrt(3 / (4\pi)) ( 1/r^3 * [-yx, x^2+z^2, -yz]  ) (m=-1)
        \nabla Y00  = \sqrt(3 / (4\pi)) ( 1/r^3 * [-zx, -zy, x^2+y^2]  ) (m= 0)
        \nabla Y01  = \sqrt(3 / (4\pi)) ( 1/r^3 * [y^2+z^2, -xy, -xz]  ) (m=-1)
    """

    r = torch.sqrt((xyz**2).sum(3))
    r3 = r**3
    c = 0.4886025119029199
    p = (c / r3).unsqueeze(-1)

    if m == -1:
        return p * (torch.stack([-xyz[:, :, :, 1] * xyz[:, :, :, 0],
                                 xyz[:, :, :, 0]**2 + xyz[:, :, :, 2]**2,
                                 -xyz[:, :, :, 1] * xyz[:, :, :, 2]],
                                dim=-1))
    if m == 0:

        return p * (torch.stack([-xyz[:, :, :, 2] * xyz[:, :, :, 0],
                                 -xyz[:, :, :, 2] * xyz[:, :, :, 1],
                                 xyz[:, :, :, 0]**2 + xyz[:, :, :, 1]**2],
                                dim=-1))
    if m == 1:
        return p * (torch.stack([xyz[:, :, :, 1]**2 + xyz[:, :, :, 2]**2,
                                 -xyz[:, :, :, 0] * xyz[:, :, :, 1],
                                 -xyz[:, :, :, 0] * xyz[:, :, :, 2]],
                                dim=-1))