/**
 * This file includes code originally from the DROID-SLAM repository:
 * https://github.com/princeton-vl/DROID-SLAM
 * Licensed under the BSD-3 License. See THIRD_PARTY_LICENSES.md for details.
 */

#include <torch/extension.h>
#include <vector>

#define CHECK_CONTIGUOUS(x) TORCH_CHECK(x.is_contiguous(), #x " must be contiguous")
#define CHECK_INPUT(x) CHECK_CONTIGUOUS(x)

std::vector<torch::Tensor> corr_index_cuda_forward(torch::Tensor volume, torch::Tensor coords, int radius);
std::vector<torch::Tensor> corr_index_cuda_backward(torch::Tensor volume, torch::Tensor coords, torch::Tensor corr_grad,
                                                    int radius);
std::vector<torch::Tensor> altcorr_cuda_forward(torch::Tensor fmap1, torch::Tensor fmap2, torch::Tensor coords,
                                                int radius);
std::vector<torch::Tensor> altcorr_cuda_backward(torch::Tensor fmap1, torch::Tensor fmap2, torch::Tensor coords,
                                                 torch::Tensor corr_grad, int radius);

// c++ python binding
std::vector<torch::Tensor> corr_index_forward(torch::Tensor volume, torch::Tensor coords, int radius) {
    CHECK_INPUT(volume);
    CHECK_INPUT(coords);

    return corr_index_cuda_forward(volume, coords, radius);
}

std::vector<torch::Tensor> corr_index_backward(torch::Tensor volume, torch::Tensor coords, torch::Tensor corr_grad,
                                               int radius) {
    CHECK_INPUT(volume);
    CHECK_INPUT(coords);
    CHECK_INPUT(corr_grad);

    auto volume_grad = corr_index_cuda_backward(volume, coords, corr_grad, radius);
    return {volume_grad};
}

std::vector<torch::Tensor> altcorr_forward(torch::Tensor fmap1, torch::Tensor fmap2, torch::Tensor coords, int radius) {
    CHECK_INPUT(fmap1);
    CHECK_INPUT(fmap2);
    CHECK_INPUT(coords);

    return altcorr_cuda_forward(fmap1, fmap2, coords, radius);
}

std::vector<torch::Tensor> altcorr_backward(torch::Tensor fmap1, torch::Tensor fmap2, torch::Tensor coords,
                                            torch::Tensor corr_grad, int radius) {
    CHECK_INPUT(fmap1);
    CHECK_INPUT(fmap2);
    CHECK_INPUT(coords);
    CHECK_INPUT(corr_grad);

    return altcorr_cuda_backward(fmap1, fmap2, coords, corr_grad, radius);
}

void pybind_droid_net_ext(py::module& m) {
    // correlation volume kernels
    m.def("altcorr_forward", &altcorr_forward, "ALTCORR forward");
    m.def("altcorr_backward", &altcorr_backward, "ALTCORR backward");
    m.def("corr_index_forward", &corr_index_forward, "INDEX forward");
    m.def("corr_index_backward", &corr_index_backward, "INDEX backward");
}