#ifndef CAFFE2_OPERATORS_UTILS_NMS_H_
#define CAFFE2_OPERATORS_UTILS_NMS_H_

#include <vector>

#include "caffe2/core/logging.h"
#include "caffe2/core/macros.h"
#include "caffe2/utils/eigen_utils.h"
#include "caffe2/utils/math.h"

namespace caffe2 {
namespace utils {

// Greedy non-maximum suppression for proposed bounding boxes
// Reject a bounding box if its region has an intersection-overunion (IoU)
//    overlap with a higher scoring selected bounding box larger than a
//    threshold.
// Reference: facebookresearch/Detectron/detectron/utils/cython_nms.pyx
// proposals: pixel coordinates of proposed bounding boxes,
//    size: (M, 4), format: [x1; y1; x2; y2]
// scores: scores for each bounding box, size: (M, 1)
// sorted_indices: indices that sorts the scores from high to low
// return: row indices of the selected proposals
template <class Derived1, class Derived2>
std::vector<int> nms_cpu_upright(
    const Eigen::ArrayBase<Derived1>& proposals,
    const Eigen::ArrayBase<Derived2>& scores,
    const std::vector<int>& sorted_indices,
    float thresh,
    int topN = -1,
    bool legacy_plus_one = false) {
  CAFFE_ENFORCE_EQ(proposals.rows(), scores.rows());
  CAFFE_ENFORCE_EQ(proposals.cols(), 4);
  CAFFE_ENFORCE_EQ(scores.cols(), 1);
  CAFFE_ENFORCE_LE(sorted_indices.size(), proposals.rows());

  using EArrX = EArrXt<typename Derived1::Scalar>;

  auto x1 = proposals.col(0);
  auto y1 = proposals.col(1);
  auto x2 = proposals.col(2);
  auto y2 = proposals.col(3);

  EArrX areas =
      (x2 - x1 + int(legacy_plus_one)) * (y2 - y1 + int(legacy_plus_one));

  EArrXi order = AsEArrXt(sorted_indices);
  std::vector<int> keep;
  while (order.size() > 0) {
    // exit if already enough proposals
    if (topN >= 0 && keep.size() >= topN) {
      break;
    }

    int i = order[0];
    keep.push_back(i);
    ConstEigenVectorArrayMap<int> rest_indices(
        order.data() + 1, order.size() - 1);
    EArrX xx1 = GetSubArray(x1, rest_indices).cwiseMax(x1[i]);
    EArrX yy1 = GetSubArray(y1, rest_indices).cwiseMax(y1[i]);
    EArrX xx2 = GetSubArray(x2, rest_indices).cwiseMin(x2[i]);
    EArrX yy2 = GetSubArray(y2, rest_indices).cwiseMin(y2[i]);

    EArrX w = (xx2 - xx1 + int(legacy_plus_one)).cwiseMax(0.0);
    EArrX h = (yy2 - yy1 + int(legacy_plus_one)).cwiseMax(0.0);
    EArrX inter = w * h;
    EArrX ovr = inter / (areas[i] + GetSubArray(areas, rest_indices) - inter);

    // indices for sub array order[1:n]
    auto inds = GetArrayIndices(ovr <= thresh);
    order = GetSubArray(order, AsEArrXt(inds) + 1);
  }

  return keep;
}

/**
 * Soft-NMS implementation as outlined in https://arxiv.org/abs/1704.04503.
 * Reference: facebookresearch/Detectron/detectron/utils/cython_nms.pyx
 * out_scores: Output updated scores after applying Soft-NMS
 * proposals: pixel coordinates of proposed bounding boxes,
 *    size: (M, 4), format: [x1; y1; x2; y2]
 *    size: (M, 5), format: [ctr_x; ctr_y; w; h; angle (degrees)] for RRPN
 * scores: scores for each bounding box, size: (M, 1)
 * indices: Indices to consider within proposals and scores. Can be used
 *     to pre-filter proposals/scores based on some threshold.
 * sigma: Standard deviation for Gaussian
 * overlap_thresh: Similar to original NMS
 * score_thresh: If updated score falls below this thresh, discard proposal
 * method: 0 - Hard (original) NMS, 1 - Linear, 2 - Gaussian
 * return: row indices of the selected proposals
 */
template <class Derived1, class Derived2, class Derived3>
std::vector<int> soft_nms_cpu_upright(
    Eigen::ArrayBase<Derived3>* out_scores,
    const Eigen::ArrayBase<Derived1>& proposals,
    const Eigen::ArrayBase<Derived2>& scores,
    const std::vector<int>& indices,
    float sigma = 0.5,
    float overlap_thresh = 0.3,
    float score_thresh = 0.001,
    unsigned int method = 1,
    int topN = -1,
    bool legacy_plus_one = false) {
  CAFFE_ENFORCE_EQ(proposals.rows(), scores.rows());
  CAFFE_ENFORCE_EQ(proposals.cols(), 4);
  CAFFE_ENFORCE_EQ(scores.cols(), 1);

  using EArrX = EArrXt<typename Derived1::Scalar>;

  const auto& x1 = proposals.col(0);
  const auto& y1 = proposals.col(1);
  const auto& x2 = proposals.col(2);
  const auto& y2 = proposals.col(3);

  EArrX areas =
      (x2 - x1 + int(legacy_plus_one)) * (y2 - y1 + int(legacy_plus_one));

  // Initialize out_scores with original scores. Will be iteratively updated
  // as Soft-NMS is applied.
  *out_scores = scores;

  std::vector<int> keep;
  EArrXi pending = AsEArrXt(indices);
  while (pending.size() > 0) {
    // Exit if already enough proposals
    if (topN >= 0 && keep.size() >= topN) {
      break;
    }

    // Find proposal with max score among remaining proposals
    int max_pos;
    auto max_score = GetSubArray(*out_scores, pending).maxCoeff(&max_pos);
    int i = pending[max_pos];
    keep.push_back(i);

    // Compute IoU of the remaining boxes with the identified max box
    std::swap(pending(0), pending(max_pos));
    const auto& rest_indices = pending.tail(pending.size() - 1);
    EArrX xx1 = GetSubArray(x1, rest_indices).cwiseMax(x1[i]);
    EArrX yy1 = GetSubArray(y1, rest_indices).cwiseMax(y1[i]);
    EArrX xx2 = GetSubArray(x2, rest_indices).cwiseMin(x2[i]);
    EArrX yy2 = GetSubArray(y2, rest_indices).cwiseMin(y2[i]);

    EArrX w = (xx2 - xx1 + int(legacy_plus_one)).cwiseMax(0.0);
    EArrX h = (yy2 - yy1 + int(legacy_plus_one)).cwiseMax(0.0);
    EArrX inter = w * h;
    EArrX ovr = inter / (areas[i] + GetSubArray(areas, rest_indices) - inter);

    // Update scores based on computed IoU, overlap threshold and NMS method
    for (int j = 0; j < rest_indices.size(); ++j) {
      typename Derived2::Scalar weight;
      switch (method) {
        case 1: // Linear
          weight = (ovr(j) > overlap_thresh) ? (1.0 - ovr(j)) : 1.0;
          break;
        case 2: // Gaussian
          weight = std::exp(-1.0 * ovr(j) * ovr(j) / sigma);
          break;
        default: // Original NMS
          weight = (ovr(j) > overlap_thresh) ? 0.0 : 1.0;
      }
      (*out_scores)(rest_indices[j]) *= weight;
    }

    // Discard boxes with new scores below min threshold and update pending
    // indices
    const auto& rest_scores = GetSubArray(*out_scores, rest_indices);
    const auto& inds = GetArrayIndices(rest_scores >= score_thresh);
    pending = GetSubArray(rest_indices, AsEArrXt(inds));
  }

  return keep;
}

namespace {
const int INTERSECT_NONE = 0;
const int INTERSECT_PARTIAL = 1;
const int INTERSECT_FULL = 2;

class RotatedRect {
 public:
  RotatedRect() {}
  RotatedRect(
      const Eigen::Vector2f& p_center,
      const Eigen::Vector2f& p_size,
      float p_angle)
      : center(p_center), size(p_size), angle(p_angle) {}
  void get_vertices(Eigen::Vector2f* pt) const {
    // M_PI / 180. == 0.01745329251
    double _angle = angle * 0.01745329251;
    float b = (float)cos(_angle) * 0.5f;
    float a = (float)sin(_angle) * 0.5f;

    pt[0].x() = center.x() - a * size.y() - b * size.x();
    pt[0].y() = center.y() + b * size.y() - a * size.x();
    pt[1].x() = center.x() + a * size.y() - b * size.x();
    pt[1].y() = center.y() - b * size.y() - a * size.x();
    pt[2] = 2 * center - pt[0];
    pt[3] = 2 * center - pt[1];
  }
  Eigen::Vector2f center;
  Eigen::Vector2f size;
  float angle;
};

template <class Derived>
RotatedRect bbox_to_rotated_rect(const Eigen::ArrayBase<Derived>& box) {
  CAFFE_ENFORCE_EQ(box.size(), 5);
  // cv::RotatedRect takes angle to mean clockwise rotation, but RRPN bbox
  // representation means counter-clockwise rotation.
  return RotatedRect(
      Eigen::Vector2f(box[0], box[1]),
      Eigen::Vector2f(box[2], box[3]),
      -box[4]);
}

// Eigen doesn't seem to support 2d cross product, so we make one here
float cross_2d(const Eigen::Vector2f& A, const Eigen::Vector2f& B) {
  return A.x() * B.y() - B.x() * A.y();
}

// rotated_rect_intersection_pts is a replacement function for
// cv::rotatedRectangleIntersection, which has a bug due to float underflow
// For anyone interested, here're the PRs on OpenCV:
// https://github.com/opencv/opencv/issues/12221
// https://github.com/opencv/opencv/pull/12222
// Note that we do not check if the number of intersections is <= 8 in this case
int rotated_rect_intersection_pts(
    const RotatedRect& rect1,
    const RotatedRect& rect2,
    Eigen::Vector2f* intersections,
    int& num) {
  // Used to test if two points are the same
  const float samePointEps = 0.00001f;
  const float EPS = 1e-14;
  num = 0; // number of intersections

  Eigen::Vector2f vec1[4], vec2[4], pts1[4], pts2[4];

  rect1.get_vertices(pts1);
  rect2.get_vertices(pts2);

  // Specical case of rect1 == rect2
  bool same = true;

  for (int i = 0; i < 4; i++) {
    if (fabs(pts1[i].x() - pts2[i].x()) > samePointEps ||
        (fabs(pts1[i].y() - pts2[i].y()) > samePointEps)) {
      same = false;
      break;
    }
  }

  if (same) {
    for (int i = 0; i < 4; i++) {
      intersections[i] = pts1[i];
    }
    num = 4;
    return INTERSECT_FULL;
  }

  // Line vector
  // A line from p1 to p2 is: p1 + (p2-p1)*t, t=[0,1]
  for (int i = 0; i < 4; i++) {
    vec1[i] = pts1[(i + 1) % 4] - pts1[i];
    vec2[i] = pts2[(i + 1) % 4] - pts2[i];
  }

  // Line test - test all line combos for intersection
  for (int i = 0; i < 4; i++) {
    for (int j = 0; j < 4; j++) {
      // Solve for 2x2 Ax=b

      // This takes care of parallel lines
      float det = cross_2d(vec2[j], vec1[i]);
      if (std::fabs(det) <= EPS) {
        continue;
      }

      auto vec12 = pts2[j] - pts1[i];

      float t1 = cross_2d(vec2[j], vec12) / det;
      float t2 = cross_2d(vec1[i], vec12) / det;

      if (t1 >= 0.0f && t1 <= 1.0f && t2 >= 0.0f && t2 <= 1.0f) {
        intersections[num++] = pts1[i] + t1 * vec1[i];
      }
    }
  }

  // Check for vertices from rect1 inside rect2
  {
    const auto& AB = vec2[0];
    const auto& DA = vec2[3];
    auto ABdotAB = AB.squaredNorm();
    auto ADdotAD = DA.squaredNorm();
    for (int i = 0; i < 4; i++) {
      // assume ABCD is the rectangle, and P is the point to be judged
      // P is inside ABCD iff. P's projection on AB lies within AB
      // and P's projection on AD lies within AD

      auto AP = pts1[i] - pts2[0];

      auto APdotAB = AP.dot(AB);
      auto APdotAD = -AP.dot(DA);

      if ((APdotAB >= 0) && (APdotAD >= 0) && (APdotAB <= ABdotAB) &&
          (APdotAD <= ADdotAD)) {
        intersections[num++] = pts1[i];
      }
    }
  }

  // Reverse the check - check for vertices from rect2 inside rect1
  {
    const auto& AB = vec1[0];
    const auto& DA = vec1[3];
    auto ABdotAB = AB.squaredNorm();
    auto ADdotAD = DA.squaredNorm();
    for (int i = 0; i < 4; i++) {
      auto AP = pts2[i] - pts1[0];

      auto APdotAB = AP.dot(AB);
      auto APdotAD = -AP.dot(DA);

      if ((APdotAB >= 0) && (APdotAD >= 0) && (APdotAB <= ABdotAB) &&
          (APdotAD <= ADdotAD)) {
        intersections[num++] = pts2[i];
      }
    }
  }

  return num ? INTERSECT_PARTIAL : INTERSECT_NONE;
}

// Compute convex hull using Graham scan algorithm
int convex_hull_graham(
    const Eigen::Vector2f* p,
    const int& num_in,
    Eigen::Vector2f* q,
    bool shift_to_zero = false) {
  CAFFE_ENFORCE(num_in >= 2);
  std::vector<int> order;