import torch


def knn_point(k, points):
    """Returns kNN idx for given pointcloud.

    Args:
        k (int): kNN neighbor num
        points (tensor): batch pc, [B, K, 3]
    """
    # no grad
    pc = points.clone().detach()
    # build kNN graph
    B, K = pc.shape[:2]
    pc = pc.transpose(2, 1)  # [B, 3, K]
    inner = -2. * torch.matmul(pc.transpose(2, 1), pc)  # [B, K, K]
    xx = torch.sum(pc ** 2, dim=1, keepdim=True)  # [B, 1, K]
    dist = xx + inner + xx.transpose(2, 1)  # [B, K, K], l2^2
    assert dist.min().item() >= -1e-4
    # the min is self so we take top (k + 1)
    _, top_idx = (-dist).topk(k=k + 1, dim=-1)  # [B, K, k + 1]
    top_idx = top_idx[:, :, 1:]  # [B, K, k]
    return top_idx