#include <ATen/cuda/CUDAContext.h>
#include <torch/extension.h>
// NOTE: AT_CHECK has become TORCH_CHECK
void gather_points_kernel_wrapper(int b, int c, int n, int npoints,
                                  const float *points, const int *idx,
                                  float *out);
void gather_points_grad_kernel_wrapper(int b, int c, int n, int npoints,
                                       const float *grad_out, const int *idx,
                                       float *grad_points);

void minimum_density_sampling_kernel_wrapper(int b, int n, int m,
                                             const float *dataset, float *temp,
                                             int *idxs, float *mean_mst_length);

#define CUDA_CHECK_ERRORS()                                           \
  do                                                                  \
  {                                                                   \
    cudaError_t err = cudaGetLastError();                             \
    if (cudaSuccess != err)                                           \
    {                                                                 \
      fprintf(stderr, "CUDA kernel failed : %s\n%s at L:%d in %s\n",  \
              cudaGetErrorString(err), __PRETTY_FUNCTION__, __LINE__, \
              __FILE__);                                              \
      exit(-1);                                                       \
    }                                                                 \
  } while (0)

#define CHECK_CUDA(x)                                             \
  do                                                              \
  {                                                               \
    TORCH_CHECK(x.type().is_cuda(), #x " must be a CUDA tensor"); \
  } while (0)

#define CHECK_CONTIGUOUS(x)                                            \
  do                                                                   \
  {                                                                    \
    TORCH_CHECK(x.is_contiguous(), #x " must be a contiguous tensor"); \
  } while (0)

#define CHECK_IS_INT(x)                                 \
  do                                                    \
  {                                                     \
    TORCH_CHECK(x.scalar_type() == at::ScalarType::Int, \
                #x " must be an int tensor");           \
  } while (0)

#define CHECK_IS_FLOAT(x)                                 \
  do                                                      \
  {                                                       \
    TORCH_CHECK(x.scalar_type() == at::ScalarType::Float, \
                #x " must be a float tensor");            \
  } while (0)

at::Tensor gather_points(at::Tensor points, at::Tensor idx)
{
  CHECK_CONTIGUOUS(points);
  CHECK_CONTIGUOUS(idx);
  CHECK_IS_FLOAT(points);
  CHECK_IS_INT(idx);

  if (points.type().is_cuda())
  {
    CHECK_CUDA(idx);
  }

  at::Tensor output =
      torch::zeros({points.size(0), points.size(1), idx.size(1)},
                   at::device(points.device()).dtype(at::ScalarType::Float));

  if (points.type().is_cuda())
  {
    gather_points_kernel_wrapper(points.size(0), points.size(1), points.size(2),
                                 idx.size(1), points.data<float>(),
                                 idx.data<int>(), output.data<float>());
  }
  else
  {
    TORCH_CHECK(false, "CPU not supported");
  }

  return output;
}

at::Tensor gather_points_grad(at::Tensor grad_out, at::Tensor idx,
                              const int n)
{
  CHECK_CONTIGUOUS(grad_out);
  CHECK_CONTIGUOUS(idx);
  CHECK_IS_FLOAT(grad_out);
  CHECK_IS_INT(idx);

  if (grad_out.type().is_cuda())
  {
    CHECK_CUDA(idx);
  }

  at::Tensor output =
      torch::zeros({grad_out.size(0), grad_out.size(1), n},
                   at::device(grad_out.device()).dtype(at::ScalarType::Float));

  if (grad_out.type().is_cuda())
  {
    gather_points_grad_kernel_wrapper(grad_out.size(0), grad_out.size(1), n,
                                      idx.size(1), grad_out.data<float>(),
                                      idx.data<int>(), output.data<float>());
  }
  else
  {
    TORCH_CHECK(false, "CPU not supported");
  }

  return output;
}
at::Tensor minimum_density_sampling(at::Tensor points, const int nsamples, at::Tensor mean_mst_length, at::Tensor output)
{
  CHECK_CONTIGUOUS(points);
  CHECK_IS_FLOAT(points);

  at::Tensor tmp =
      torch::zeros({points.size(0), points.size(1)},
                   at::device(points.device()).dtype(at::ScalarType::Float));

  if (points.type().is_cuda())
  {
    minimum_density_sampling_kernel_wrapper(
        points.size(0), points.size(1), nsamples, points.data<float>(),
        tmp.data<float>(), output.data<int>(), mean_mst_length.data<float>());
  }
  else
  {
    TORCH_CHECK(false, "CPU not supported");
  }

  return output;
}

PYBIND11_MODULE(TORCH_EXTENSION_NAME, m)
{
  m.def("minimum_density_sampling", &minimum_density_sampling, "minimum_density_sampling (CUDA)");
  m.def("gather_points_grad", &gather_points_grad, "gather_points_grad (CUDA)");
  m.def("gather_points", &gather_points, "gather_points (CUDA)");
}
