#include <torch/library.h>

#include <ATen/cuda/CUDAContext.h>
#include <c10/cuda/CUDAGuard.h>

#include "../vox_cuda/vox_kernels.h"
#include "registration.h"
#include "torch_binding.h"

namespace {

void check_dims(int64_t dx, int64_t dy, int64_t dz) {
  TORCH_CHECK(dx >= 1 && dy >= 1 && dz >= 1 && dx <= 8192 && dy <= 8192 &&
                  dz <= 8192,
              "grid dims must be in [1, 8192]");
}

void check_grid(const at::Tensor& grid, int64_t dx, int64_t dy, int64_t dz) {
  const int64_t words = ((dz + 31) / 32) * dy * dx;
  TORCH_CHECK(grid.is_cuda() && grid.scalar_type() == at::kInt &&
                  grid.is_contiguous() && grid.numel() == words,
              "grid must be a contiguous CUDA int32 tensor with "
              "ceil(dz/32)*dy*dx elements");
}

void check_mesh(const at::Tensor& verts, const at::Tensor& faces,
                const at::Tensor& tris) {
  TORCH_CHECK(verts.is_cuda() && verts.scalar_type() == at::kFloat &&
                  verts.dim() == 2 && verts.size(1) == 3 &&
                  verts.is_contiguous(),
              "verts must be a contiguous CUDA float32 [V, 3] tensor");
  TORCH_CHECK(faces.is_cuda() && faces.scalar_type() == at::kInt &&
                  faces.dim() == 2 && faces.size(1) == 3 &&
                  faces.is_contiguous(),
              "faces must be a contiguous CUDA int32 [F, 3] tensor");
  TORCH_CHECK(tris.is_cuda() && tris.scalar_type() == at::kInt &&
                  tris.dim() == 1 && tris.is_contiguous(),
              "tris must be a contiguous CUDA int32 1-D tensor");
}

void check_points(const at::Tensor& points) {
  TORCH_CHECK(points.is_cuda() && points.scalar_type() == at::kFloat &&
                  points.dim() == 2 && points.size(1) == 3 &&
                  points.is_contiguous(),
              "points must be a contiguous CUDA float32 [N, 3] tensor");
}

}  // namespace

void vox_surface(at::Tensor verts, at::Tensor faces, at::Tensor tris,
                 at::Tensor grid, int64_t dx, int64_t dy, int64_t dz,
                 double ox, double oy, double oz, double inv_vs,
                 int64_t group) {
  check_dims(dx, dy, dz);
  check_mesh(verts, faces, tris);
  check_grid(grid, dx, dy, dz);
  TORCH_CHECK(group == 1 || group == 32 || group == 256,
              "group must be 1, 32, or 256");
  const c10::cuda::CUDAGuard guard(grid.device());
  auto stream = at::cuda::getCurrentCUDAStream();
  const cudaError_t err = vox_surface_launch(
      verts.data_ptr<float>(), faces.data_ptr<int>(), tris.data_ptr<int>(),
      tris.numel(), reinterpret_cast<unsigned int*>(grid.data_ptr<int>()),
      (int)dx, (int)dy, (int)dz, (float)ox, (float)oy, (float)oz,
      (float)inv_vs, (int)group, stream.stream());
  TORCH_CHECK(err == cudaSuccess, "vox_surface launch failed: ",
              cudaGetErrorString(err));
}

void vox_solid_cross(at::Tensor verts, at::Tensor faces, at::Tensor tris,
                     at::Tensor grid, int64_t dx, int64_t dy, int64_t dz,
                     double ox, double oy, double oz, double inv_vs,
                     int64_t group) {
  check_dims(dx, dy, dz);
  check_mesh(verts, faces, tris);
  check_grid(grid, dx, dy, dz);
  TORCH_CHECK(group == 1 || group == 32 || group == 256,
              "group must be 1, 32, or 256");
  const c10::cuda::CUDAGuard guard(grid.device());
  auto stream = at::cuda::getCurrentCUDAStream();
  const cudaError_t err = vox_solid_cross_launch(
      verts.data_ptr<float>(), faces.data_ptr<int>(), tris.data_ptr<int>(),
      tris.numel(), reinterpret_cast<unsigned int*>(grid.data_ptr<int>()),
      (int)dx, (int)dy, (int)dz, (float)ox, (float)oy, (float)oz,
      (float)inv_vs, (int)group, stream.stream());
  TORCH_CHECK(err == cudaSuccess, "vox_solid_cross launch failed: ",
              cudaGetErrorString(err));
}

void vox_parity(at::Tensor grid, int64_t dx, int64_t dy, int64_t dz) {
  check_dims(dx, dy, dz);
  check_grid(grid, dx, dy, dz);
  const c10::cuda::CUDAGuard guard(grid.device());
  auto stream = at::cuda::getCurrentCUDAStream();
  const cudaError_t err = vox_parity_launch(
      reinterpret_cast<unsigned int*>(grid.data_ptr<int>()), (int)dx, (int)dy,
      (int)dz, stream.stream());
  TORCH_CHECK(err == cudaSuccess, "vox_parity launch failed: ",
              cudaGetErrorString(err));
}

void vox_point_keys(at::Tensor points, at::Tensor keys, int64_t dx,
                    int64_t dy, int64_t dz, double ox, double oy, double oz,
                    double inv_vs) {
  check_dims(dx, dy, dz);
  check_points(points);
  TORCH_CHECK(keys.is_cuda() && keys.scalar_type() == at::kLong &&
                  keys.is_contiguous() && keys.numel() == points.size(0),
              "keys must be a contiguous CUDA int64 tensor with one element "
              "per point");
  const c10::cuda::CUDAGuard guard(points.device());
  auto stream = at::cuda::getCurrentCUDAStream();
  const cudaError_t err = vox_point_keys_launch(
      points.data_ptr<float>(), points.size(0),
      reinterpret_cast<long long*>(keys.data_ptr<int64_t>()), (int)dx,
      (int)dy, (int)dz, (float)ox, (float)oy, (float)oz, (float)inv_vs,
      stream.stream());
  TORCH_CHECK(err == cudaSuccess, "vox_point_keys launch failed: ",
              cudaGetErrorString(err));
}

void vox_point_occ(at::Tensor points, at::Tensor grid, int64_t dx, int64_t dy,
                   int64_t dz, double ox, double oy, double oz,
                   double inv_vs) {
  check_dims(dx, dy, dz);
  check_points(points);
  check_grid(grid, dx, dy, dz);
  const c10::cuda::CUDAGuard guard(grid.device());
  auto stream = at::cuda::getCurrentCUDAStream();
  const cudaError_t err = vox_point_occ_launch(
      points.data_ptr<float>(), points.size(0),
      reinterpret_cast<unsigned int*>(grid.data_ptr<int>()), (int)dx, (int)dy,
      (int)dz, (float)ox, (float)oy, (float)oz, (float)inv_vs,
      stream.stream());
  TORCH_CHECK(err == cudaSuccess, "vox_point_occ launch failed: ",
              cudaGetErrorString(err));
}

void vox_seg_reduce(at::Tensor feats, at::Tensor order, at::Tensor segoff,
                    at::Tensor out, int64_t mode) {
  TORCH_CHECK(feats.is_cuda() && feats.scalar_type() == at::kFloat &&
                  feats.dim() == 2 && feats.is_contiguous(),
              "feats must be a contiguous CUDA float32 [N, C] tensor");
  TORCH_CHECK(order.is_cuda() && order.scalar_type() == at::kLong &&
                  order.dim() == 1 && order.is_contiguous(),
              "order must be a contiguous CUDA int64 1-D tensor");
  TORCH_CHECK(segoff.is_cuda() && segoff.scalar_type() == at::kLong &&
                  segoff.dim() == 1 && segoff.is_contiguous() &&
                  segoff.numel() >= 1,
              "segoff must be a contiguous CUDA int64 [S+1] tensor");
  const int64_t nseg = segoff.numel() - 1;
  TORCH_CHECK(out.is_cuda() && out.scalar_type() == at::kFloat &&
                  out.dim() == 2 && out.is_contiguous() &&
                  out.size(0) == nseg && out.size(1) == feats.size(1),
              "out must be a contiguous CUDA float32 [S, C] tensor");
  TORCH_CHECK(mode >= 0 && mode <= 2, "mode must be 0 (mean), 1 (max), or "
              "2 (sum)");
  const c10::cuda::CUDAGuard guard(feats.device());
  auto stream = at::cuda::getCurrentCUDAStream();
  const cudaError_t err = vox_seg_reduce_launch(
      feats.data_ptr<float>(),
      reinterpret_cast<const long long*>(order.data_ptr<int64_t>()),
      reinterpret_cast<const long long*>(segoff.data_ptr<int64_t>()), nseg,
      feats.size(1), out.data_ptr<float>(), (int)mode, stream.stream());
  TORCH_CHECK(err == cudaSuccess, "vox_seg_reduce launch failed: ",
              cudaGetErrorString(err));
}

TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
  ops.def("vox_surface(Tensor verts, Tensor faces, Tensor tris, Tensor grid, "
          "int dx, int dy, int dz, float ox, float oy, float oz, "
          "float inv_vs, int group) -> ()");
  ops.def("vox_solid_cross(Tensor verts, Tensor faces, Tensor tris, "
          "Tensor grid, int dx, int dy, int dz, float ox, float oy, "
          "float oz, float inv_vs, int group) -> ()");
  ops.def("vox_parity(Tensor grid, int dx, int dy, int dz) -> ()");
  ops.def("vox_point_keys(Tensor points, Tensor keys, int dx, int dy, "
          "int dz, float ox, float oy, float oz, float inv_vs) -> ()");
  ops.def("vox_point_occ(Tensor points, Tensor grid, int dx, int dy, int dz, "
          "float ox, float oy, float oz, float inv_vs) -> ()");
  ops.def("vox_seg_reduce(Tensor feats, Tensor order, Tensor segoff, "
          "Tensor out, int mode) -> ()");

#if defined(CUDA_KERNEL) || defined(ROCM_KERNEL)
  ops.impl("vox_surface", torch::kCUDA, &vox_surface);
  ops.impl("vox_solid_cross", torch::kCUDA, &vox_solid_cross);
  ops.impl("vox_parity", torch::kCUDA, &vox_parity);
  ops.impl("vox_point_keys", torch::kCUDA, &vox_point_keys);
  ops.impl("vox_point_occ", torch::kCUDA, &vox_point_occ);
  ops.impl("vox_seg_reduce", torch::kCUDA, &vox_seg_reduce);
#endif
}

REGISTER_EXTENSION(TORCH_EXTENSION_NAME)