#include <torch/library.h>
#include <torch/all.h>

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

#include "registration.h"
#include "torch_binding.h"
#include "../pathtracer_diff_cuda/geometry_grad.h"
#include "../pathtracer_diff_cuda/pathtracer_launch.h"

namespace {

const float* fptr(const torch::Tensor& t) {
  return t.numel() ? t.const_data_ptr<float>() : nullptr;
}
const int* iptr(const torch::Tensor& t) {
  return t.numel() ? t.const_data_ptr<int>() : nullptr;
}

void chk(bool ok, const char* msg) { TORCH_CHECK(ok, msg); }

PtdSceneArgs pack_args(
    const torch::Tensor& tris, const torch::Tensor& mat_ids,
    const torch::Tensor& uvs, const torch::Tensor& nodes_f,
    const torch::Tensor& nodes_i, const torch::Tensor& light_faces,
    const torch::Tensor& light_cdf, double total_light_area,
    const torch::Tensor& tex, const torch::Tensor& tex_hdr,
    const torch::Tensor& emi_tex, const torch::Tensor& emi_hdr,
    const torch::Tensor& mat_type, const torch::Tensor& mat_rough,
    const torch::Tensor& mat_ior, const torch::Tensor& med_sa,
    const torch::Tensor& med_ss, double med_sbar, const torch::Tensor& env,
    const torch::Tensor& env_cdf_m, const torch::Tensor& env_cdf_c,
    const torch::Tensor& env_pdf, int64_t env_w, int64_t env_h, int64_t spp,
    int64_t max_bounces, int64_t mode) {
  chk(tris.is_cuda() && tris.is_contiguous() &&
          tris.dtype() == torch::kFloat32 && tris.dim() == 2 &&
          tris.size(1) == 9,
      "tris must be contiguous CUDA f32 [F, 9]");
  chk(mat_ids.dtype() == torch::kInt32 && mat_ids.numel() == tris.size(0),
      "mat_ids must be i32 [F]");
  chk(uvs.numel() == tris.size(0) * 6, "uvs must be [F, 3, 2]");
  chk(nodes_f.dim() == 2 && nodes_f.size(1) == 6, "nodes_f [N, 6]");
  chk(nodes_i.dim() == 2 && nodes_i.size(1) == 3 &&
          nodes_i.size(0) == nodes_f.size(0),
      "nodes_i [N, 3]");
  chk(light_cdf.numel() == light_faces.numel(), "light list mismatch");
  chk(tex.dim() == 2 && tex.size(1) == 3, "tex [T, 3]");
  chk(tex_hdr.dim() == 2 && tex_hdr.size(1) == 3, "tex_hdr [M, 3]");
  int64_t M = tex_hdr.size(0);
  chk(M <= 64, "at most 64 materials");
  chk(emi_tex.dim() == 2 && emi_tex.size(1) == 3, "emi_tex [Te, 3]");
  chk(emi_hdr.sizes() == tex_hdr.sizes(), "emi_hdr [M, 3]");
  chk(mat_type.numel() == M && mat_rough.numel() == M && mat_ior.numel() == M,
      "per-material arrays must have M entries");
  chk(med_sa.numel() == med_ss.numel() &&
          (med_sa.numel() == 0 || med_sa.numel() == 3),
      "medium sigmas must be [3] or empty");
  if (env.numel()) {
    chk(env.numel() == env_w * env_h * 3 && env_cdf_m.numel() == env_h &&
            env_cdf_c.numel() == env_w * env_h &&
            env_pdf.numel() == env_w * env_h,
        "env tables mismatch");
  }
  chk(spp >= 1, "spp must be >= 1");
  chk(max_bounces >= 1 && max_bounces <= 16, "max_bounces in [1, 16]");
  chk(mode >= 0 && mode <= 2, "mode must be 0|1|2");

  PtdSceneArgs a;
  a.tris = tris.const_data_ptr<float>();
  a.mat_ids = mat_ids.const_data_ptr<int>();
  a.uvs = uvs.const_data_ptr<float>();
  a.n_faces = (int)tris.size(0);
  a.nodes_f = nodes_f.const_data_ptr<float>();
  a.nodes_i = nodes_i.const_data_ptr<int>();
  a.n_nodes = (int)nodes_f.size(0);
  a.light_faces = iptr(light_faces);
  a.light_cdf = fptr(light_cdf);
  a.n_lights = (int)light_faces.numel();
  a.total_light_area = (float)total_light_area;
  a.tex = tex.const_data_ptr<float>();
  a.tex_hdr = tex_hdr.const_data_ptr<int>();
  a.n_texels = (int)tex.size(0);
  a.emi_tex = emi_tex.const_data_ptr<float>();
  a.emi_hdr = emi_hdr.const_data_ptr<int>();
  a.n_emi_texels = (int)emi_tex.size(0);
  a.mat_type = mat_type.const_data_ptr<int>();
  a.mat_rough = mat_rough.const_data_ptr<float>();
  a.mat_ior = mat_ior.const_data_ptr<float>();
  a.n_mats = (int)M;
  a.med_sa = fptr(med_sa);
  a.med_ss = fptr(med_ss);
  a.med_sbar = (float)med_sbar;
  a.has_med = med_sa.numel() ? 1 : 0;
  a.env = fptr(env);
  a.env_w = (int)env_w;
  a.env_h = (int)env_h;
  a.env_cdf_m = fptr(env_cdf_m);
  a.env_cdf_c = fptr(env_cdf_c);
  a.env_pdf = fptr(env_pdf);
  return a;
}

}  // namespace

void pt_forward(torch::Tensor tris, torch::Tensor mat_ids, torch::Tensor uvs,
                torch::Tensor nodes_f, torch::Tensor nodes_i,
                torch::Tensor light_faces, torch::Tensor light_cdf,
                double total_light_area, torch::Tensor tex,
                torch::Tensor tex_hdr, torch::Tensor emi_tex,
                torch::Tensor emi_hdr, torch::Tensor mat_type,
                torch::Tensor mat_rough, torch::Tensor mat_ior,
                torch::Tensor med_sa, torch::Tensor med_ss, double med_sbar,
                torch::Tensor env, torch::Tensor env_cdf_m,
                torch::Tensor env_cdf_c, torch::Tensor env_pdf, int64_t env_w,
                int64_t env_h, torch::Tensor cam, int64_t spp,
                int64_t max_bounces, int64_t mode, int64_t seed,
                torch::Tensor image) {
  PtdSceneArgs a = pack_args(tris, mat_ids, uvs, nodes_f, nodes_i, light_faces,
                             light_cdf, total_light_area, tex, tex_hdr,
                             emi_tex, emi_hdr, mat_type, mat_rough, mat_ior,
                             med_sa, med_ss, med_sbar, env, env_cdf_m,
                             env_cdf_c, env_pdf, env_w, env_h, spp,
                             max_bounces, mode);
  TORCH_CHECK(image.is_cuda() && image.is_contiguous() &&
                  image.dtype() == torch::kFloat32 && image.dim() == 3 &&
                  image.size(2) == 3,
              "image must be contiguous CUDA f32 [H, W, 3]");
  const at::cuda::CUDAGuard guard(tris.device());
  cudaStream_t stream = at::cuda::getCurrentCUDAStream();
  torch::Tensor cam_h = cam.to(torch::kFloat32).to(torch::kCPU).contiguous();
  ptd_forward_launch(&a, cam_h.const_data_ptr<float>(), (int)image.size(0),
                     (int)image.size(1), (int)spp, (int)max_bounces,
                     (int)mode, (long long)seed, image.data_ptr<float>(),
                     stream);
  C10_CUDA_KERNEL_LAUNCH_CHECK();
}

void pt_backward(torch::Tensor tris, torch::Tensor mat_ids, torch::Tensor uvs,
                 torch::Tensor nodes_f, torch::Tensor nodes_i,
                 torch::Tensor light_faces, torch::Tensor light_cdf,
                 double total_light_area, torch::Tensor tex,
                 torch::Tensor tex_hdr, torch::Tensor emi_tex,
                 torch::Tensor emi_hdr, torch::Tensor mat_type,
                 torch::Tensor mat_rough, torch::Tensor mat_ior,
                 torch::Tensor med_sa, torch::Tensor med_ss, double med_sbar,
                 torch::Tensor env, torch::Tensor env_cdf_m,
                 torch::Tensor env_cdf_c, torch::Tensor env_pdf, int64_t env_w,
                 int64_t env_h, torch::Tensor cam, int64_t spp,
                 int64_t max_bounces, int64_t mode, int64_t seed,
                 torch::Tensor grad_image, torch::Tensor grad_tex,
                 torch::Tensor grad_emi_tex, torch::Tensor grad_env,
                 torch::Tensor grad_med) {
  PtdSceneArgs a = pack_args(tris, mat_ids, uvs, nodes_f, nodes_i, light_faces,
                             light_cdf, total_light_area, tex, tex_hdr,
                             emi_tex, emi_hdr, mat_type, mat_rough, mat_ior,
                             med_sa, med_ss, med_sbar, env, env_cdf_m,
                             env_cdf_c, env_pdf, env_w, env_h, spp,
                             max_bounces, mode);
  TORCH_CHECK(grad_image.is_cuda() && grad_image.is_contiguous() &&
                  grad_image.dim() == 3 && grad_image.size(2) == 3,
              "grad_image must be contiguous CUDA f32 [H, W, 3]");
  TORCH_CHECK(grad_tex.sizes() == tex.sizes(), "grad_tex must match tex");
  TORCH_CHECK(grad_emi_tex.sizes() == emi_tex.sizes(),
              "grad_emi_tex must match emi_tex");
  TORCH_CHECK(grad_env.sizes() == env.sizes(), "grad_env must match env");
  TORCH_CHECK(grad_med.numel() == (med_sa.numel() ? 6 : 0),
              "grad_med must be [6] with a medium, else empty");
  const at::cuda::CUDAGuard guard(tris.device());
  cudaStream_t stream = at::cuda::getCurrentCUDAStream();
  torch::Tensor cam_h = cam.to(torch::kFloat32).to(torch::kCPU).contiguous();
  ptd_backward_launch(&a, cam_h.const_data_ptr<float>(),
                      (int)grad_image.size(0), (int)grad_image.size(1),
                      (int)spp, (int)max_bounces, (int)mode, (long long)seed,
                      grad_image.const_data_ptr<float>(),
                      grad_tex.data_ptr<float>(),
                      grad_emi_tex.data_ptr<float>(),
                      grad_env.numel() ? grad_env.data_ptr<float>() : nullptr,
                      grad_med.numel() ? grad_med.data_ptr<float>() : nullptr,
                      stream);
  C10_CUDA_KERNEL_LAUNCH_CHECK();
}

void pt_geometry_grad(torch::Tensor tris, torch::Tensor mat_ids,
                      torch::Tensor uvs, torch::Tensor nodes_f,
                      torch::Tensor nodes_i, torch::Tensor light_faces,
                      torch::Tensor light_cdf, double total_light_area,
                      torch::Tensor tex, torch::Tensor tex_hdr,
                      torch::Tensor emi_tex, torch::Tensor emi_hdr,
                      torch::Tensor mat_type, torch::Tensor mat_rough,
                      torch::Tensor mat_ior, torch::Tensor face_verts,
                      torch::Tensor edges, torch::Tensor edge_cdf,
                      torch::Tensor cam, int64_t spp, int64_t edge_samples,
                      int64_t seed, torch::Tensor grad_image,
                      torch::Tensor grad_verts) {
  torch::Tensor z = torch::zeros(0, tris.options());
  PtdSceneArgs a = pack_args(tris, mat_ids, uvs, nodes_f, nodes_i,
                             light_faces, light_cdf, total_light_area, tex,
                             tex_hdr, emi_tex, emi_hdr, mat_type, mat_rough,
                             mat_ior, z, z, 0.0, z, z, z, z, 0, 0,
                             /*spp=*/1, /*max_bounces=*/4, /*mode=*/0);
  TORCH_CHECK(face_verts.dtype() == torch::kInt32 &&
                  face_verts.numel() == tris.size(0) * 3,
              "face_verts must be i32 [F, 3]");
  TORCH_CHECK(grad_image.is_cuda() && grad_image.is_contiguous() &&
                  grad_image.dim() == 3 && grad_image.size(2) == 3,
              "grad_image must be contiguous CUDA f32 [H, W, 3]");
  TORCH_CHECK(grad_verts.dim() == 2 && grad_verts.size(1) == 3,
              "grad_verts must be [V, 3]");
  TORCH_CHECK(edges.numel() == 0 ||
                  (edges.dtype() == torch::kInt32 && edges.dim() == 2 &&
                   edges.size(1) == 4 &&
                   edge_cdf.numel() == edges.size(0) + 1),
              "edges [E, 4] i32 with edge_cdf [E + 1]");
  const at::cuda::CUDAGuard guard(tris.device());
  cudaStream_t stream = at::cuda::getCurrentCUDAStream();
  torch::Tensor cam_h = cam.to(torch::kFloat32).to(torch::kCPU).contiguous();
  int H = (int)grad_image.size(0), W = (int)grad_image.size(1);
  ptd_geo_interior_launch(&a, face_verts.const_data_ptr<int>(),
                          cam_h.const_data_ptr<float>(), H, W, (int)spp,
                          (long long)seed, grad_image.const_data_ptr<float>(),
                          grad_verts.data_ptr<float>(), stream);
  if (edges.numel())
    ptd_geo_boundary_launch(&a, face_verts.const_data_ptr<int>(),
                            edges.const_data_ptr<int>(), (int)edges.size(0),
                            edge_cdf.const_data_ptr<float>(),
                            cam_h.const_data_ptr<float>(), H, W,
                            (int)edge_samples, (long long)seed,
                            grad_image.const_data_ptr<float>(),
                            grad_verts.data_ptr<float>(), stream);
  C10_CUDA_KERNEL_LAUNCH_CHECK();
}

TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
  ops.def(
      "pt_forward(Tensor tris, Tensor mat_ids, Tensor uvs, Tensor nodes_f,"
      " Tensor nodes_i, Tensor light_faces, Tensor light_cdf,"
      " float total_light_area, Tensor tex, Tensor tex_hdr, Tensor emi_tex,"
      " Tensor emi_hdr, Tensor mat_type, Tensor mat_rough, Tensor mat_ior,"
      " Tensor med_sa, Tensor med_ss, float med_sbar, Tensor env,"
      " Tensor env_cdf_m, Tensor env_cdf_c, Tensor env_pdf, int env_w,"
      " int env_h, Tensor cam, int spp, int max_bounces, int mode, int seed,"
      " Tensor! image) -> ()");
  ops.impl("pt_forward", torch::kCUDA, &pt_forward);

  ops.def(
      "pt_backward(Tensor tris, Tensor mat_ids, Tensor uvs, Tensor nodes_f,"
      " Tensor nodes_i, Tensor light_faces, Tensor light_cdf,"
      " float total_light_area, Tensor tex, Tensor tex_hdr, Tensor emi_tex,"
      " Tensor emi_hdr, Tensor mat_type, Tensor mat_rough, Tensor mat_ior,"
      " Tensor med_sa, Tensor med_ss, float med_sbar, Tensor env,"
      " Tensor env_cdf_m, Tensor env_cdf_c, Tensor env_pdf, int env_w,"
      " int env_h, Tensor cam, int spp, int max_bounces, int mode, int seed,"
      " Tensor grad_image, Tensor! grad_tex, Tensor! grad_emi_tex,"
      " Tensor! grad_env, Tensor! grad_med) -> ()");
  ops.impl("pt_backward", torch::kCUDA, &pt_backward);

  ops.def(
      "pt_geometry_grad(Tensor tris, Tensor mat_ids, Tensor uvs,"
      " Tensor nodes_f, Tensor nodes_i, Tensor light_faces, Tensor light_cdf,"
      " float total_light_area, Tensor tex, Tensor tex_hdr, Tensor emi_tex,"
      " Tensor emi_hdr, Tensor mat_type, Tensor mat_rough, Tensor mat_ior,"
      " Tensor face_verts, Tensor edges, Tensor edge_cdf, Tensor cam,"
      " int spp, int edge_samples, int seed, Tensor grad_image,"
      " Tensor! grad_verts) -> ()");
  ops.impl("pt_geometry_grad", torch::kCUDA, &pt_geometry_grad);
}

REGISTER_EXTENSION(TORCH_EXTENSION_NAME)