#pragma once

#include <torch/torch.h>

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);

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);

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);