// Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved.
#pragma once

#ifdef WITH_CUDA
#include "cuda/vision.h"
#endif

// Interface for Python
at::Tensor ROIAlign_3d_forward(const at::Tensor& input,
                               const at::Tensor& rois,
                               const double spatial_scale,
                               const int64_t pooled_depth,
                               const int64_t pooled_height,
                               const int64_t pooled_width,
                               const int64_t sampling_ratio) {
    if (input.type().is_cuda()) {
#ifdef WITH_CUDA
        return ROIAlign_3d_forward_cuda(input, rois, spatial_scale, pooled_depth, pooled_height, pooled_width, sampling_ratio);
#else
        AT_ERROR("Not compiled with GPU support");
#endif
    }
    AT_ERROR("Not implemented on the CPU");
}

at::Tensor ROIAlign_3d_backward(const at::Tensor& grad,
                                const at::Tensor& rois,
                                const double spatial_scale,
                                const int64_t pooled_depth,
                                const int64_t pooled_height,
                                const int64_t pooled_width,
                                const int64_t batch_size,
                                const int64_t channels,
                                const int64_t depth,
                                const int64_t height,
                                const int64_t width,
                                const int64_t sampling_ratio) {
    if (grad.type().is_cuda()) {
#ifdef WITH_CUDA
        return ROIAlign_3d_backward_cuda(grad, rois, spatial_scale, pooled_depth, pooled_height, pooled_width, batch_size, channels, depth, height, width, sampling_ratio);
#else
        AT_ERROR("Not compiled with GPU support");
#endif
    }
    AT_ERROR("Not implemented on the CPU");
}