#include <ATen/cuda/CUDAContext.h>
#include <cuda_runtime.h>
#include <torch/extension.h>

#include "../include/quadrilinear4d_cuda.h"
#include "../include/quadrilinear4d_kernel.h"

int quadrilinear4d_forward_cuda(torch::Tensor lut, torch::Tensor image,
                                torch::Tensor output, int lut_dim, int shift,
                                float binsize, int width, int height,
                                int batch) {
  const float* lut_flat = lut.data_ptr<float>();
  const float* image_flat = image.data_ptr<float>();
  float* output_flat = output.data_ptr<float>();

  QuadriLinearForwardLaucher(lut_flat, image_flat, output_flat, lut_dim, shift,
                             binsize, width, height, batch,
                             at::cuda::getCurrentCUDAStream());
  return 1;
}

int quadrilinear4d_backward_cuda(torch::Tensor image, torch::Tensor image_grad,
                                 torch::Tensor lut, torch::Tensor lut_grad,
                                 int lut_dim, int shift, float binsize,
                                 int width, int height, int batch) {
  const float* image_flat = image.data_ptr<float>();
  float* image_grad_flat = image_grad.data_ptr<float>();
  const float* lut_flat = lut.data_ptr<float>();
  float* lut_grad_flat = lut_grad.data_ptr<float>();

  QuadriLinearBackwardLaucher(image_flat, image_grad_flat, lut_flat,
                              lut_grad_flat, lut_dim, shift, binsize, width,
                              height, batch, at::cuda::getCurrentCUDAStream());
  return 1;
}