// Modified from https://github.com/pytorch/pytorch/blob/master/aten/src/ATen/native/AffineGridGenerator.cpp
#include <ATen/ATen.h>
#include <ATen/NativeFunctions.h>
#include <torch/extension.h>

namespace mmdetection {

using namespace at;

at::Tensor linspace_from_neg_one(const Tensor& grid, int64_t num_steps,
                                 bool align_corners) {
  if (num_steps <= 1) {
    return at::tensor(0, grid.options());
  }
  auto range = at::linspace(-1, 1, num_steps, grid.options());
  if (!align_corners) {
    range = range * (num_steps - 1) / num_steps;
  }
  return range;
}

Tensor make_base_grid_4D(const Tensor& theta, int64_t N, int64_t C, int64_t H,
                         int64_t W, bool align_corners) {
  auto base_grid = at::empty({N, H, W, 3}, theta.options());

  base_grid.select(-1, 0).copy_(linspace_from_neg_one(theta, W, align_corners));
  base_grid.select(-1, 1).copy_(
      linspace_from_neg_one(theta, H, align_corners).unsqueeze_(-1));
  base_grid.select(-1, 2).fill_(1);

  return base_grid;
}

Tensor make_base_grid_5D(const Tensor& theta, int64_t N, int64_t C, int64_t D,
                         int64_t H, int64_t W, bool align_corners) {
  auto base_grid = at::empty({N, D, H, W, 4}, theta.options());

  base_grid.select(-1, 0).copy_(linspace_from_neg_one(theta, W, align_corners));
  base_grid.select(-1, 1).copy_(
      linspace_from_neg_one(theta, H, align_corners).unsqueeze_(-1));
  base_grid.select(-1, 2).copy_(linspace_from_neg_one(theta, D, align_corners)
                                    .unsqueeze_(-1)
                                    .unsqueeze_(-1));
  base_grid.select(-1, 3).fill_(1);

  return base_grid;
}

Tensor affine_grid_generator_4D_forward(const Tensor& theta, int64_t N,
                                        int64_t C, int64_t H, int64_t W,
                                        bool align_corners) {
  Tensor base_grid = make_base_grid_4D(theta, N, C, H, W, align_corners);
  auto grid = base_grid.view({N, H * W, 3}).bmm(theta.transpose(1, 2));
  return grid.view({N, H, W, 2});
}

Tensor affine_grid_generator_5D_forward(const Tensor& theta, int64_t N,
                                        int64_t C, int64_t D, int64_t H,
                                        int64_t W, bool align_corners) {
  Tensor base_grid = make_base_grid_5D(theta, N, C, D, H, W, align_corners);
  auto grid = base_grid.view({N, D * H * W, 4}).bmm(theta.transpose(1, 2));
  return grid.view({N, D, H, W, 3});
}

Tensor affine_grid_generator_forward(const Tensor& theta, IntArrayRef size,
                                     bool align_corners) {
  if (size.size() == 4) {
    return affine_grid_generator_4D_forward(theta, size[0], size[1], size[2],
                                            size[3], align_corners);
  } else {
    return affine_grid_generator_5D_forward(theta, size[0], size[1], size[2],
                                            size[3], size[4], align_corners);
  }
}

Tensor affine_grid_generator_4D_backward(const Tensor& grad_grid, int64_t N,
                                         int64_t C, int64_t H, int64_t W,
                                         bool align_corners) {
  auto base_grid = make_base_grid_4D(grad_grid, N, C, H, W, align_corners);
  AT_ASSERT(grad_grid.sizes() == IntArrayRef({N, H, W, 2}));
  auto grad_theta = base_grid.view({N, H * W, 3})
                        .transpose(1, 2)
                        .bmm(grad_grid.view({N, H * W, 2}));
  return grad_theta.transpose(1, 2);
}

Tensor affine_grid_generator_5D_backward(const Tensor& grad_grid, int64_t N,
                                         int64_t C, int64_t D, int64_t H,
                                         int64_t W, bool align_corners) {
  auto base_grid = make_base_grid_5D(grad_grid, N, C, D, H, W, align_corners);
  AT_ASSERT(grad_grid.sizes() == IntArrayRef({N, D, H, W, 3}));
  auto grad_theta = base_grid.view({N, D * H * W, 4})
                        .transpose(1, 2)
                        .bmm(grad_grid.view({N, D * H * W, 3}));
  return grad_theta.transpose(1, 2);
}

Tensor affine_grid_generator_backward(const Tensor& grad, IntArrayRef size,
                                      bool align_corners) {
  if (size.size() == 4) {
    return affine_grid_generator_4D_backward(grad, size[0], size[1], size[2],
                                             size[3], align_corners);
  } else {
    return affine_grid_generator_5D_backward(grad, size[0], size[1], size[2],
                                             size[3], size[4], align_corners);
  }
}
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
  m.def("affine_grid_generator_forward", &affine_grid_generator_forward,
        "affine_grid_generator_forward");
  m.def("affine_grid_generator_backward", &affine_grid_generator_backward,
        "affine_grid_generator_backward");
}

}  // namespace mmdetection
