/* 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. */
#pragma once
#include <map>
#include "paddle/framework/selected_rows.h"
#include "paddle/platform/device_context.h"

namespace paddle {
namespace operators {
namespace math {
template <typename T>
struct BBox {
  BBox(T x_min, T y_min, T x_max, T y_max)
      : x_min(x_min),
        y_min(y_min),
        x_max(x_max),
        y_max(y_max),
        is_difficult(false) {}

  BBox() {}

  T get_width() const { return x_max - x_min; }

  T get_height() const { return y_max - y_min; }

  T get_center_x() const { return (x_min + x_max) / 2; }

  T get_center_y() const { return (y_min + y_max) / 2; }

  T get_area() const { return get_width() * get_height(); }

  // coordinate of bounding box
  T x_min;
  T y_min;
  T x_max;
  T y_max;
  // whether difficult object (e.g. object with heavy occlusion is difficult)
  bool is_difficult;
};
// KNCHW ==> NHWC
// template <typename T>
template <typename T>
void get_bbox_from_priorData(const T* prior_data, const size_t num_bboxes,
                             std::vector<BBox<T>>& bbox_vec);
template <typename T>
void get_bbox_var_from_prior_data(const T* prior_data, const size_t num,
                                  std::vector<std::vector<T>>& var_vec);
template <typename T>
BBox<T> decode_bbox_with_var(BBox<T>& prior_bbox,
                             const std::vector<T>& prior_bbox_var,
                             const std::vector<T>& loc_pred_data);
template <typename T1, typename T2>
bool sort_score_pair_descend(const std::pair<T1, T2>& pair1,
                             const std::pair<T1, T2>& pair2);
template <typename T>
bool sort_score_pair_descend(const std::pair<T, BBox<T>>& pair1,
                             const std::pair<T, BBox<T>>& pair2);
template <typename T>
T jaccard_overlap(const BBox<T>& bbox1, const BBox<T>& bbox2);

template <typename T>
void apply_nms_fast(const std::vector<BBox<T>>& bboxes,
                    const T* conf_score_data, size_t class_idx, size_t top_k,
                    T conf_threshold, T nms_threshold, size_t num_priors,
                    size_t num_classes, std::vector<size_t>* indices);
template <typename T>
int get_detection_indices(
    const T* conf_data, const size_t num_priors, const size_t num_classes,
    const size_t background_label_id, const size_t batch_size,
    const T conf_threshold, const size_t nms_top_k, const T nms_threshold,
    const size_t top_k,
    const std::vector<std::vector<BBox<T>>>& all_decoded_bboxes,
    std::vector<std::map<size_t, std::vector<size_t>>>* all_detection_indices);
template <typename T>
BBox<T> clipBBox(const BBox<T>& bbox);
template <typename T>
void get_detection_output(
    const T* conf_data, const size_t num_kept, const size_t num_priors,
    const size_t num_classes, const size_t batch_size,
    const std::vector<std::map<size_t, std::vector<size_t>>>& all_indices,
    const std::vector<std::vector<BBox<T>>>& all_decoded_bboxes, T* out_data);
template <typename T>
void get_bbox_from_priorData(const T* prior_data, const size_t num_bboxes,
                             std::vector<BBox<T>>& bbox_vec) {
  size_t out_offset = bbox_vec.size();
  bbox_vec.resize(bbox_vec.size() + num_bboxes);
  for (size_t i = 0; i < num_bboxes; ++i) {
    BBox<T> bbox;
    bbox.x_min = *(prior_data + i * 8);
    bbox.y_min = *(prior_data + i * 8 + 1);
    bbox.x_max = *(prior_data + i * 8 + 2);
    bbox.y_max = *(prior_data + i * 8 + 3);
    bbox_vec[out_offset + i] = bbox;
  }
}
template <typename T>
void get_bbox_var_from_prior_data(const T* prior_data, const size_t num,
                                  std::vector<std::vector<T>>& var_vec) {
  size_t out_offset = var_vec.size();
  var_vec.resize(var_vec.size() + num);
  for (size_t i = 0; i < num; ++i) {
    std::vector<T> var;
    var.push_back(*(prior_data + i * 8 + 4));
    var.push_back(*(prior_data + i * 8 + 5));
    var.push_back(*(prior_data + i * 8 + 6));
    var.push_back(*(prior_data + i * 8 + 7));
    var_vec[out_offset + i] = var;
  }
}
template <typename T>
BBox<T> decode_bbox_with_var(BBox<T>& prior_bbox,
                             const std::vector<T>& prior_bbox_var,
                             const std::vector<T>& loc_pred_data) {
  T prior_bbox_width = prior_bbox.get_width();
  T prior_bbox_height = prior_bbox.get_height();
  T prior_bbox_center_x = prior_bbox.get_center_x();
  T prior_bbox_center_y = prior_bbox.get_center_y();

  T decoded_bbox_center_x =
      prior_bbox_var[0] * loc_pred_data[0] * prior_bbox_width +
      prior_bbox_center_x;
  T decoded_bbox_center_y =
      prior_bbox_var[1] * loc_pred_data[1] * prior_bbox_height +
      prior_bbox_center_y;
  T decoded_bbox_width =
      std::exp(prior_bbox_var[2] * loc_pred_data[2]) * prior_bbox_width;
  T decoded_bbox_height =
      std::exp(prior_bbox_var[3] * loc_pred_data[3]) * prior_bbox_height;

  BBox<T> decoded_bbox;
  decoded_bbox.x_min = decoded_bbox_center_x - decoded_bbox_width / 2;
  decoded_bbox.y_min = decoded_bbox_center_y - decoded_bbox_height / 2;
  decoded_bbox.x_max = decoded_bbox_center_x + decoded_bbox_width / 2;
  decoded_bbox.y_max = decoded_bbox_center_y + decoded_bbox_height / 2;

  return decoded_bbox;
}
template <typename T1, typename T2>
bool sort_score_pair_descend(const std::pair<T1, T2>& pair1,
                             const std::pair<T1, T2>& pair2) {
  return pair1.first > pair2.first;
}
template <typename T>
T jaccard_overlap(const BBox<T>& bbox1, const BBox<T>& bbox2) {
  if (bbox2.x_min > bbox1.x_max || bbox2.x_max < bbox1.x_min ||
      bbox2.y_min > bbox1.y_max || bbox2.y_max < bbox1.y_min) {
    return 0.0;
  } else {
    T inter_x_min = std::max(bbox1.x_min, bbox2.x_min);
    T inter_y_min = std::max(bbox1.y_min, bbox2.y_min);
    T interX_max = std::min(bbox1.x_max, bbox2.x_max);
    T interY_max = std::min(bbox1.y_max, bbox2.y_max);

    T inter_width = interX_max - inter_x_min;
    T inter_height = interY_max - inter_y_min;
    T inter_area = inter_width * inter_height;

    T bbox_area1 = bbox1.get_area();
    T bbox_area2 = bbox2.get_area();

    return inter_area / (bbox_area1 + bbox_area2 - inter_area);
  }
}

template <typename T>
void apply_nms_fast(const std::vector<BBox<T>>& bboxes,
                    const T* conf_score_data, size_t class_idx, size_t top_k,
                    T conf_threshold, T nms_threshold, size_t num_priors,
                    size_t num_classes, std::vector<size_t>* indices) {
  std::vector<std::pair<T, size_t>> scores;
  for (size_t i = 0; i < num_priors; ++i) {
    size_t conf_offset = i * num_classes + class_idx;
    if (conf_score_data[conf_offset] > conf_threshold)
      scores.push_back(std::make_pair(conf_score_data[conf_offset], i));
  }
  std::stable_sort(scores.begin(), scores.end(),
                   sort_score_pair_descend<T, size_t>);
  if (top_k > 0 && top_k < scores.size()) scores.resize(top_k);
  while (scores.size() > 0) {
    const size_t idx = scores.front().second;
    bool keep = true;
    for (size_t i = 0; i < indices->size(); ++i) {
      if (keep) {
        const size_t saved_idx = (*indices)[i];
        T overlap = jaccard_overlap<T>(bboxes[idx], bboxes[saved_idx]);
        keep = overlap <= nms_threshold;
      } else {
        break;
      }
    }
    if (keep) indices->push_back(idx);
    scores.erase(scores.begin());
  }
}
template <typename T>
int get_detection_indices(
    const T* conf_data, const size_t num_priors, const size_t num_classes,
    const size_t background_label_id, const size_t batch_size,
    const T conf_threshold, const size_t nms_top_k, const T nms_threshold,
    const size_t top_k,
    const std::vector<std::vector<BBox<T>>>& all_decoded_bboxes,
    std::vector<std::map<size_t, std::vector<size_t>>>* all_detection_indices) {
  int total_keep_num = 0;
  for (size_t n = 0; n < batch_size; ++n) {
    const std::vector<BBox<T>>& decoded_bboxes = all_decoded_bboxes[n];
    size_t num_detected = 0;
    std::map<size_t, std::vector<size_t>> indices;
    size_t conf_offset = n * num_priors * num_classes;
    for (size_t c = 0; c < num_classes; ++c) {
      if (c == background_label_id) continue;
      apply_nms_fast<T>(decoded_bboxes, conf_data + conf_offset, c, nms_top_k,
                        conf_threshold, nms_threshold, num_priors, num_classes,
                        &(indices[c]));
      num_detected += indices[c].size();
    }
    if (top_k > 0 && num_detected > top_k) {
      // std::vector<pair<T,T>> score_index_pairs;
      std::vector<std::pair<T, std::pair<size_t, size_t>>> score_index_pairs;
      for (size_t c = 0; c < num_classes; ++c) {
        const std::vector<size_t>& label_indices = indices[c];
        for (size_t i = 0; i < label_indices.size(); ++i) {
          size_t idx = label_indices[i];
          score_index_pairs.push_back(
              std::make_pair((conf_data + conf_offset)[idx * num_classes + c],
                             std::make_pair(c, idx)));
        }
      }
      std::sort(score_index_pairs.begin(), score_index_pairs.end(),
                sort_score_pair_descend<T, std::pair<size_t, size_t>>);
      score_index_pairs.resize(top_k);
      std::map<size_t, std::vector<size_t>> new_indices;
      for (size_t i = 0; i < score_index_pairs.size(); ++i) {
        size_t label = score_index_pairs[i].second.first;
        size_t idx = score_index_pairs[i].second.second;
        new_indices[label].push_back(idx);
      }
      all_detection_indices->push_back(new_indices);
      total_keep_num += top_k;
    } else {
      all_detection_indices->push_back(indices);
      total_keep_num += num_detected;
    }
  }
  return total_keep_num;
}
template <typename T>
BBox<T> clipBBox(const BBox<T>& bbox) {
  T one = static_cast<T>(1.0);
  T zero = static_cast<T>(0.0);
  BBox<T> clipped_bbox;
  clipped_bbox.x_min = std::max(std::min(bbox.x_min, one), zero);
  clipped_bbox.y_min = std::max(std::min(bbox.y_min, one), zero);
  clipped_bbox.x_max = std::max(std::min(bbox.x_max, one), zero);
  clipped_bbox.y_max = std::max(std::min(bbox.y_max, one), zero);
  return clipped_bbox;
}
template <typename T>
void get_detection_output(
    const T* conf_data, const size_t num_kept, const size_t num_priors,
    const size_t num_classes, const size_t batch_size,
    const std::vector<std::map<size_t, std::vector<size_t>>>& all_indices,
    const std::vector<std::vector<BBox<T>>>& all_decoded_bboxes, T* out_data) {
  size_t count = 0;
  for (size_t n = 0; n < batch_size; ++n) {
    for (std::map<size_t, std::vector<size_t>>::const_iterator it =
             all_indices[n].begin();
         it != all_indices[n].end(); ++it) {
      size_t label = it->first;
      const std::vector<size_t>& indices = it->second;
      const std::vector<BBox<T>>& decoded_bboxes = all_decoded_bboxes[n];
      for (size_t i = 0; i < indices.size(); ++i) {
        size_t idx = indices[i];
        size_t conf_offset = n * num_priors * num_classes + idx * num_classes;
        out_data[count * 7] = n;
        out_data[count * 7 + 1] = label;
        out_data[count * 7 + 2] = (conf_data + conf_offset)[label];
        BBox<T> clipped_bbox = clipBBox<T>(decoded_bboxes[idx]);
        out_data[count * 7 + 3] = clipped_bbox.x_min;
        out_data[count * 7 + 4] = clipped_bbox.y_min;
        out_data[count * 7 + 5] = clipped_bbox.x_max;
        out_data[count * 7 + 6] = clipped_bbox.y_max;
        ++count;
      }
    }
  }
}
}  // namespace math
}  // namespace operators
}  // namespace paddle