/* Copyright (c) 2016 PaddlePaddle Authors. All Rights Reserve.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License. */

#include "paddle/framework/op_registry.h"
#include "paddle/operators/math/math_function.h"

namespace paddle {
namespace operators {

using Tensor = framework::Tensor;
using LoDTensor = framework::LoDTensor;

constexpr int64_t kOutputDim = 6;
constexpr int64_t kBBoxSize = 4;

class MulticlassNMSOp : public framework::OperatorWithKernel {
 public:
  using framework::OperatorWithKernel::OperatorWithKernel;

  void InferShape(framework::InferShapeContext* ctx) const override {
    PADDLE_ENFORCE(ctx->HasInput("Bboxes"),
                   "Input(Bboxes) of MulticlassNMS should not be null.");
    PADDLE_ENFORCE(ctx->HasInput("Scores"),
                   "Input(Scores) of MulticlassNMS should not be null.");

    auto box_dims = ctx->GetInputDim("Bboxes");
    auto score_dims = ctx->GetInputDim("Scores");

    PADDLE_ENFORCE_EQ(box_dims.size(), 3,
                      "The rank of Input(Bboxes) must be 3.");
    PADDLE_ENFORCE_EQ(score_dims.size(), 3,
                      "The rank of Input(Scores) must be 3.");
    PADDLE_ENFORCE_EQ(box_dims[0], score_dims[0]);
    PADDLE_ENFORCE_EQ(box_dims[2], 4);
    PADDLE_ENFORCE_EQ(box_dims[1], score_dims[2]);

    // Here the box_dims[0] is not the real dimension of output.
    // It will be rewritten in the computing kernel.
    ctx->SetOutputDim("Out", {box_dims[0], 6});
  }
};

template <class T>
bool SortScorePairDescend(const std::pair<float, T>& pair1,
                          const std::pair<float, T>& pair2) {
  return pair1.first > pair2.first;
}

template <class T>
static inline void GetMaxScoreIndex(
    const std::vector<T>& scores, const T threshold, int top_k,
    std::vector<std::pair<T, int>>* sorted_indices) {
  for (size_t i = 0; i < scores.size(); ++i) {
    if (scores[i] > threshold) {
      sorted_indices->push_back(std::make_pair(scores[i], i));
    }
  }
  // Sort the score pair according to the scores in descending order
  std::stable_sort(sorted_indices->begin(), sorted_indices->end(),
                   SortScorePairDescend<int>);
  // Keep top_k scores if needed.
  if (top_k > -1 && top_k < sorted_indices->size()) {
    sorted_indices->resize(top_k);
  }
}

template <class T>
T BBoxArea(const T* box, const bool normalized) {
  if (box[2] < box[0] || box[3] < box[1]) {
    // If bbox is invalid (e.g. xmax < xmin or ymax < ymin), return 0.
    return T(0.);
  } else {
    const T w = box[2] - box[0];
    const T h = box[3] - box[1];
    if (normalized) {
      return w * h;
    } else {
      // If bbox is not within range [0, 1].
      return (w + 1) * (h + 1);
    }
  }
}

template <class T>
static inline T JaccardOverlap(const T* box1, const T* box2,
                               const bool normalized) {
  if (box2[0] > box1[2] || box2[2] < box1[0] || box2[1] > box1[3] ||
      box2[3] < box1[1]) {
    return static_cast<T>(0.);
  } else {
    const T inter_xmin = std::max(box1[0], box2[0]);
    const T inter_ymin = std::max(box1[1], box2[1]);
    const T inter_xmax = std::min(box1[2], box2[2]);
    const T inter_ymax = std::min(box1[3], box2[3]);
    const T inter_w = inter_xmax - inter_xmin;
    const T inter_h = inter_ymax - inter_ymin;
    const T inter_area = inter_w * inter_h;
    const T bbox1_area = BBoxArea<T>(box1, normalized);
    const T bbox2_area = BBoxArea<T>(box2, normalized);
    return inter_area / (bbox1_area + bbox2_area - inter_area);
  }
}

template <typename T>
class MulticlassNMSKernel : public framework::OpKernel<T> {
 public:
  void NMSFast(const Tensor& bbox, const Tensor& scores,
               const T score_threshold, const T nms_threshold, const T eta,
               const int64_t top_k, std::vector<int>* selected_indices) const {
    // The total boxes for each instance.
    int64_t num_boxes = bbox.dims()[0];
    // 4: [xmin ymin xmax ymax]
    int64_t box_size = bbox.dims()[1];

    std::vector<T> scores_data(num_boxes);
    std::copy_n(scores.data<T>(), num_boxes, scores_data.begin());
    std::vector<std::pair<T, int>> sorted_indices;
    GetMaxScoreIndex(scores_data, score_threshold, top_k, &sorted_indices);

    selected_indices->clear();
    T adaptive_threshold = nms_threshold;
    const T* bbox_data = bbox.data<T>();

    while (sorted_indices.size() != 0) {
      const int idx = sorted_indices.front().second;
      bool keep = true;
      for (int k = 0; k < selected_indices->size(); ++k) {
        if (keep) {
          const int kept_idx = (*selected_indices)[k];
          T overlap = JaccardOverlap<T>(bbox_data + idx * box_size,
                                        bbox_data + kept_idx * box_size, true);
          keep = overlap <= adaptive_threshold;
        } else {
          break;
        }
      }
      if (keep) {
        selected_indices->push_back(idx);
      }
      sorted_indices.erase(sorted_indices.begin());
      if (keep && eta < 1 && adaptive_threshold > 0.5) {
        adaptive_threshold *= eta;
      }
    }
  }

  void MulticlassNMS(const framework::ExecutionContext& ctx,
                     const Tensor& scores, const Tensor& bboxes,
                     std::map<int, std::vector<int>>* indices,
                     int* num_nmsed_out) const {
    int64_t background_label = ctx.Attr<int64_t>("background_label");
    int64_t nms_top_k = ctx.Attr<int64_t>("nms_top_k");
    int64_t keep_top_k = ctx.Attr<int64_t>("keep_top_k");
    T nms_threshold = static_cast<T>(ctx.Attr<float>("nms_threshold"));
    T nms_eta = static_cast<T>(ctx.Attr<float>("nms_eta"));
    T score_threshold = static_cast<T>(ctx.Attr<float>("confidence_threshold"));

    int64_t class_num = scores.dims()[0];
    int64_t predict_dim = scores.dims()[1];
    int num_det = 0;
    for (int64_t c = 0; c < class_num; ++c) {
      if (c == background_label) continue;
      Tensor score = scores.Slice(c, c + 1);
      NMSFast(bboxes, score, score_threshold, nms_threshold, nms_eta, nms_top_k,
              &((*indices)[c]));
      num_det += indices[c].size();
    }

    *num_nmsed_out = num_det;
    const T* scores_data = scores.data<T>();
    if (keep_top_k > -1 && num_det > keep_top_k) {
      std::vector<std::pair<float, std::pair<int, int>>> score_index_pairs;
      for (const auto& it : *indices) {
        int label = it.first;
        const T* sdata = scores_data + label * predict_dim;
        const std::vector<int>& label_indices = it.second;
        for (int j = 0; j < label_indices.size(); ++j) {
          int idx = label_indices[j];
          PADDLE_ENFORCE_LT(idx, predict_dim);
          score_index_pairs.push_back(
              std::make_pair(sdata[idx], std::make_pair(label, idx)));
        }
      }
      // Keep top k results per image.
      std::sort(score_index_pairs.begin(), score_index_pairs.end(),
                SortScorePairDescend<std::pair<int, int>>);
      score_index_pairs.resize(keep_top_k);

      // Store the new indices.
      std::map<int, std::vector<int>> new_indices;
      for (int j = 0; j < score_index_pairs.size(); ++j) {
        int label = score_index_pairs[j].second.first;
        int idx = score_index_pairs[j].second.second;
        new_indices[label].push_back(idx);
      }
      new_indices.swap(*indices);
      *num_nmsed_out = keep_top_k;
    }
  }

  void MulticlassOutput(const Tensor& scores, const Tensor& bboxes,
                        std::map<int, std::vector<int>>& selected_indices,
                        Tensor* outs) const {
    int predict_dim = scores.dims()[1];
    auto* scores_data = scores.data<T>();
    auto* bboxes_data = bboxes.data<T>();
    auto* odata = outs->data<T>();

    int count = 0;
    for (const auto& it : selected_indices) {
      int label = it.first;
      const T* sdata = scores_data + label * predict_dim;
      std::vector<int> indices = it.second;
      for (int j = 0; j < indices.size(); ++j) {
        int idx = indices[j];
        const T* bdata = bboxes_data + idx * kBBoxSize;
        odata[count * kOutputDim] = label;           // label
        odata[count * kOutputDim + 1] = sdata[idx];  // score
        odata[count * kOutputDim + 2] = bdata[0];    // xmin
        odata[count * kOutputDim + 3] = bdata[1];    // ymin
        odata[count * kOutputDim + 4] = bdata[2];    // xmax
        odata[count * kOutputDim + 5] = bdata[3];    // ymax
      }
      count++;
    }
  }

  void Compute(const framework::ExecutionContext& ctx) const override {
    auto* boxes = ctx.Input<Tensor>("Bboxes");
    auto* scores = ctx.Input<Tensor>("Scores");
    auto* outs = ctx.Output<LoDTensor>("Out");

    auto box_dims = boxes->dims();
    auto score_dims = scores->dims();

    int64_t batch_size = box_dims[0];
    int64_t class_num = score_dims[1];
    int64_t predict_dim = score_dims[2];

    std::vector<std::map<int, std::vector<int>>> all_indices;
    std::vector<size_t> batch_starts = {0};
    for (int64_t i = 0; i < batch_size; ++i) {
      Tensor ins_score = scores->Slice(i, i + 1);
      ins_score.Resize({class_num, predict_dim});
      std::map<int, std::vector<int>> indices;
      int num_nmsed_out = 0;
      MulticlassNMS(ctx, ins_score, *boxes, &indices, &num_nmsed_out);
      all_indices.push_back(indices);
      batch_starts.push_back(batch_starts.back() + num_nmsed_out);
    }

    int num_kept = batch_starts.back();
    if (num_kept == 0) {
      outs->Resize({0, 0});
    } else {
      outs->mutable_data<T>({num_kept, kOutputDim}, ctx.GetPlace());
      for (int64_t i = 0; i < batch_size; ++i) {
        Tensor ins_score = scores->Slice(i, i + 1);
        ins_score.Resize({class_num, predict_dim});
        int64_t s = batch_starts[i];
        int64_t e = batch_starts[i + 1];
        if (e > s) {
          Tensor out = outs->Slice(s, e);
          MulticlassOutput(ins_score, *boxes, all_indices[i], &out);
        }
      }
    }

    framework::LoD lod;
    lod.emplace_back(batch_starts);

    outs->set_lod(lod);
  }
};

class MulticlassNMSOpMaker : public framework::OpProtoAndCheckerMaker {
 public:
  MulticlassNMSOpMaker(OpProto* proto, OpAttrChecker* op_checker)
      : OpProtoAndCheckerMaker(proto, op_checker) {
    AddInput("Bboxes",
             "(Tensor) A 2-D Tensor with shape [M, 4] represents the location "
             "predictions with M bboxes. 4 is the number of "
             "each location coordinates.");
    AddOutput("Scores",
              "(Tensor) A 3-D Tensor with shape [N, C, M] represents the "
              "confidence predictions. N is the batch size, C is the class "
              "number, M is number of predictions for each class, which is "
              "the same with Bboxes.");
    AddAttr<int64_t>(
        "background_label",
        "(int64_t, defalut: 0) "
        "The index of background label, the background label will be ignored.")
        .SetDefault(0);
    AddAttr<float>("nms_threshold",
                   "(float, defalut: 0.3) "
                   "The threshold to be used in nms.")
        .SetDefault(0.3);
    AddAttr<int64_t>("nms_top_k",
                     "(int64_t) "
                     " .");
    AddAttr<float>("nms_eta",
                   "(float) "
                   "The parameter for adaptive nms.")
        .SetDefault(1.0);
    AddAttr<int64_t>("keep_top_k",
                     "(int64_t) "
                     ".");
    AddAttr<float>("confidence_threshold",
                   "(float) "
                   ".");
    AddOutput("Out",
              "(LoDTensor) A 2-D LoDTensor with shape [No, 6] represents the "
              "detections. Each row has 6 values: "
              "[label, confidence, xmin, ymin, xmax, ymax], No is the total "
              "number of detections in this mini-batch. For each instance, "
              "the offsets in first dimension are called LoD, the number of "
              "offset is N + 1, if LoD[i + 1] - LoD[i] == 0, means there is "
              "no detected bbox.");
    AddComment(R"DOC(
This operators is to do multi-class non maximum suppression (nms) on a batched
of boxes and scores.

This op greedily selects a subset of detection bounding boxes, pruning
away boxes that have high IOU (intersection over union) overlap (> thresh)
with already selected boxes.  It operates independently for each class for
which scores are provided (via the scores field of the input box_list),
pruning boxes with score less than a provided threshold prior to
applying NMS.

)DOC");
  }
};

}  // namespace operators
}  // namespace paddle

namespace ops = paddle::operators;
REGISTER_OPERATOR(multiclass_nms, ops::MulticlassNMSOp,
                  ops::MulticlassNMSOpMaker,
                  paddle::framework::EmptyGradOpMaker);
REGISTER_OP_CPU_KERNEL(multiclass_nms, ops::MulticlassNMSKernel<float>,
                       ops::MulticlassNMSKernel<double>);