// Copyright (c) Facebook, Inc. and its affiliates. All rights reserved.

#include <torch/extension.h>
#include <queue>
#include <tuple>

// Given a pixel coordinate 0 <= i < S, convert it to a normalized device
// coordinate in the range [-1, 1]. The NDC range is divided into S evenly-sized
// pixels, and assume that each pixel falls in the *center* of its range.
static float PixToNdc(const int i, const int S) {
  // NDC x-offset + (i * pixel_width + half_pixel_width)
  return -1 + (2 * i + 1.0f) / S;
}

std::tuple<torch::Tensor, torch::Tensor, torch::Tensor> RasterizePointsNaiveCpu(
    const torch::Tensor& points,
    const int image_size,
    const float radius,
    const int points_per_pixel) {
  const int N = points.size(0);
  const int P = points.size(1);
  const int S = image_size;
  const int K = points_per_pixel;
  auto int_opts = points.options().dtype(torch::kInt32);
  auto float_opts = points.options().dtype(torch::kFloat32);
  torch::Tensor point_idxs = torch::full({N, S, S, K}, -1, int_opts);
  torch::Tensor zbuf = torch::full({N, S, S, K}, -1, float_opts);
  torch::Tensor pix_dists = torch::full({N, S, S, K}, -1, float_opts);

  auto points_a = points.accessor<float, 3>();
  auto point_idxs_a = point_idxs.accessor<int32_t, 4>();
  auto zbuf_a = zbuf.accessor<float, 4>();
  auto pix_dists_a = pix_dists.accessor<float, 4>();

  const float radius2 = radius * radius;
  for (int n = 0; n < N; ++n) {
    for (int yi = 0; yi < S; ++yi) {
      float yf = PixToNdc(yi, S);
      for (int xi = 0; xi < S; ++xi) {
        float xf = PixToNdc(xi, S);
        // Use a priority queue to hold (z, idx, r)
        std::priority_queue<std::tuple<float, int, float>> q;
        for (int p = 0; p < P; ++p) {
          const float px = points_a[n][p][0];
          const float py = points_a[n][p][1];
          const float pz = points_a[n][p][2];
          if (pz < 0) {
            continue;
          }
          const float dx = px - xf;
          const float dy = py - yf;
          const float dist2 = dx * dx + dy * dy;
          if (dist2 < radius2) {
            // The current point hit the current pixel
            q.emplace(pz, p, dist2);
            if ((int)q.size() > K) {
              q.pop();
            }
          }
        }
        // Now all the points have been seen, so pop elements off the queue
        // one by one and write them into the output tensors.
        while (!q.empty()) {
          auto t = q.top();
          q.pop();
          int i = q.size();
          zbuf_a[n][yi][xi][i] = std::get<0>(t);
          point_idxs_a[n][yi][xi][i] = std::get<1>(t);
          pix_dists_a[n][yi][xi][i] = std::get<2>(t);
        }
      }
    }
  }
  return std::make_tuple(point_idxs, zbuf, pix_dists);
}

torch::Tensor RasterizePointsCoarseCpu(
    const torch::Tensor& points,
    const int image_size,
    const float radius,
    const int bin_size,
    const int max_points_per_bin) {
  const int N = points.size(0);
  const int P = points.size(1);
  const int B = 1 + (image_size - 1) / bin_size; // Integer division round up
  const int M = max_points_per_bin;
  auto opts = points.options().dtype(torch::kInt32);
  torch::Tensor points_per_bin = torch::zeros({N, B, B}, opts);
  torch::Tensor bin_points = torch::full({N, B, B, M}, -1, opts);

  auto points_a = points.accessor<float, 3>();
  auto points_per_bin_a = points_per_bin.accessor<int32_t, 3>();
  auto bin_points_a = bin_points.accessor<int32_t, 4>();

  const float pixel_width = 2.0f / image_size;
  const float bin_width = pixel_width * bin_size;
  for (int n = 0; n < N; ++n) {
    float bin_y_min = -1.0f;
    float bin_y_max = bin_y_min + bin_width;
    for (int by = 0; by < B; by++) {
      float bin_x_min = -1.0f;
      float bin_x_max = bin_x_min + bin_width;
      for (int bx = 0; bx < B; bx++) {
        int32_t points_hit = 0;
        for (int32_t p = 0; p < P; p++) {
          float px = points_a[n][p][0];
          float py = points_a[n][p][1];
          float pz = points_a[n][p][2];
          if (pz < 0) {
            continue;
          }
          float point_x_min = px - radius;
          float point_x_max = px + radius;
          float point_y_min = py - radius;
          float point_y_max = py + radius;
          // Use a half-open interval so that points exactly on the
          // boundary between bins will fall into exactly one bin.
          bool x_hit = (point_x_min <= bin_x_max) && (bin_x_min <= point_x_max);
          bool y_hit = (point_y_min <= bin_y_max) && (bin_y_min <= point_y_max);
          if (x_hit && y_hit) {
            // Got too many points for this bin, so throw an error.
            if (points_hit >= max_points_per_bin) {
              AT_ERROR("Got too many points per bin");
            }
            // The current point falls in the current bin, so
            // record it.
            bin_points_a[n][by][bx][points_hit] = p;
            points_hit++;
          }
        }
        // Record the number of points found in this bin
        points_per_bin_a[n][by][bx] = points_hit;

        // Shift the bin to the right for the next loop iteration
        bin_x_min = bin_x_max;
        bin_x_max = bin_x_min + bin_width;
      }
      // Shift the bin down for the next loop iteration
      bin_y_min = bin_y_max;
      bin_y_max = bin_y_min + bin_width;
    }
  }
  return bin_points;
}

torch::Tensor RasterizePointsBackwardCpu(
    const torch::Tensor& points, // (N, P, 3)
    const torch::Tensor& idxs, // (N, H, W, K)
    const torch::Tensor& grad_zbuf, // (N, H, W, K)
    const torch::Tensor& grad_dists) { // (N, H, W, K)
  const int N = points.size(0);
  const int P = points.size(1);
  const int H = idxs.size(1);
  const int W = idxs.size(2);
  const int K = idxs.size(3);

  // For now only support square images.
  // TODO(jcjohns): Extend to non-square images.
  if (H != W) {
    AT_ERROR("RasterizePointsBackwardCpu only supports square images");
  }
  torch::Tensor grad_points = torch::zeros({N, P, 3}, points.options());

  auto points_a = points.accessor<float, 3>();
  auto idxs_a = idxs.accessor<int32_t, 4>();
  auto grad_dists_a = grad_dists.accessor<float, 4>();
  auto grad_zbuf_a = grad_zbuf.accessor<float, 4>();
  auto grad_points_a = grad_points.accessor<float, 3>();

  for (int n = 0; n < N; ++n) { // Loop over images in the batch
    for (int y = 0; y < H; ++y) { // Loop over rows in the image
      const float yf = PixToNdc(y, H);
      for (int x = 0; x < W; ++x) { // Loop over pixels in the row
        const float xf = PixToNdc(x, W);
        for (int k = 0; k < K; ++k) { // Loop over points for the pixel
          const int p = idxs_a[n][y][x][k];
          if (p < 0) {
            break;
          }
          const float grad_dist2 = grad_dists_a[n][y][x][k];
          const float px = points_a[n][p][0];
          const float py = points_a[n][p][1];
          const float dx = px - xf;
          const float dy = py - yf;
          // Remember: dists[n][y][x][k] = dx * dx + dy * dy;
          const float grad_px = 2.0f * grad_dist2 * dx;
          const float grad_py = 2.0f * grad_dist2 * dy;
          grad_points_a[n][p][0] += grad_px;
          grad_points_a[n][p][1] += grad_py;
          grad_points_a[n][p][2] += grad_zbuf_a[n][y][x][k];
        }
      }
    }
  }
  return grad_points;
}
