// Copyright (c) OpenMMLab. All rights reserved
// Modified from
// https://github.com/sshaoshuai/Pointnet2.PyTorch/tree/master/pointnet2/src/ball_query_gpu.cu
#ifndef BALL_QUERY_CUDA_KERNEL_CUH
#define BALL_QUERY_CUDA_KERNEL_CUH

#ifdef MMCV_USE_PARROTS
#include "parrots_cuda_helper.hpp"
#else
#include "pytorch_cuda_helper.hpp"
#endif

#define DIVUP(m, n) ((m) / (n) + ((m) % (n) > 0))

template <typename T>
__global__ void ball_query_forward_cuda_kernel(int b, int n, int m,
                                               float min_radius,
                                               float max_radius, int nsample,
                                               const T* new_xyz, const T* xyz,
                                               int* idx) {
  // new_xyz: (B, M, 3)
  // xyz: (B, N, 3)
  // output:
  //      idx: (B, M, nsample)
  int bs_idx = blockIdx.y;
  int pt_idx = blockIdx.x * blockDim.x + threadIdx.x;
  if (bs_idx >= b || pt_idx >= m) return;

  new_xyz += bs_idx * m * 3 + pt_idx * 3;
  xyz += bs_idx * n * 3;
  idx += bs_idx * m * nsample + pt_idx * nsample;

  float max_radius2 = max_radius * max_radius;
  float min_radius2 = min_radius * min_radius;
  T new_x = new_xyz[0];
  T new_y = new_xyz[1];
  T new_z = new_xyz[2];

  int cnt = 0;
  for (int k = 0; k < n; ++k) {
    T x = xyz[k * 3 + 0];
    T y = xyz[k * 3 + 1];
    T z = xyz[k * 3 + 2];
    T d2 = (new_x - x) * (new_x - x) + (new_y - y) * (new_y - y) +
           (new_z - z) * (new_z - z);
    if (d2 == 0 || (d2 >= min_radius2 && d2 < max_radius2)) {
      if (cnt == 0) {
        for (int l = 0; l < nsample; ++l) {
          idx[l] = k;
        }
      }
      idx[cnt] = k;
      ++cnt;
      if (cnt >= nsample) break;
    }
  }
}

#endif  // BALL_QUERY_CUDA_KERNEL_CUH
