// Copyright (c) 2024 PaddlePaddle Authors. All Rights Reserved.
//
// 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/phi/kernels/gpu/yolo_box_post_kernel.h"
#include "paddle/common/enforce.h"
#include "paddle/phi/backends/gpu/gpu_context.h"
#include "paddle/phi/backends/gpu/gpu_launch_config.h"
#include "paddle/phi/common/memory_utils.h"
#include "paddle/phi/core/kernel_registry.h"
#include "paddle/phi/kernels/funcs/math_function.h"
#include "paddle/phi/kernels/funcs/yolo_box_util.h"
namespace phi {
struct Box {
  float x, y, w, h;
};

struct Detection {
  Box bbox;
  int classes;
  float* prob;
  float* mask;
  float objectness;
  int sort_class;
  int max_prob_class_index;
};

struct TensorInfo {
  int bbox_count_host;  // record bbox numbers
  int bbox_count_max_alloc{50};
  float* bboxes_dev_ptr;
  float* bboxes_host_ptr;
  int* bbox_count_device_ptr;  // Box counter in gpu memory, used by atomicAdd
};

static int NMSComparator(const void* pa, const void* pb) {
  const Detection a = *reinterpret_cast<const Detection*>(pa);
  const Detection b = *reinterpret_cast<const Detection*>(pb);
  if (a.max_prob_class_index > b.max_prob_class_index)
    return 1;
  else if (a.max_prob_class_index < b.max_prob_class_index)
    return -1;

  float diff = 0;
  if (b.sort_class >= 0) {
    diff = a.prob[b.sort_class] - b.prob[b.sort_class];
  } else {
    diff = a.objectness - b.objectness;
  }

  if (diff < 0)
    return 1;
  else if (diff > 0)
    return -1;
  return 0;
}

static float Overlap(float x1, float w1, float x2, float w2) {
  float l1 = x1 - w1 / 2;
  float l2 = x2 - w2 / 2;
  float left = l1 > l2 ? l1 : l2;
  float r1 = x1 + w1 / 2;
  float r2 = x2 + w2 / 2;
  float right = r1 < r2 ? r1 : r2;
  return right - left;
}

static float BoxIntersection(Box a, Box b) {
  float w = Overlap(a.x, a.w, b.x, b.w);
  float h = Overlap(a.y, a.h, b.y, b.h);
  if (w < 0 || h < 0) return 0;
  float area = w * h;
  return area;
}

static float BoxUnion(Box a, Box b) {
  float i = BoxIntersection(a, b);
  float u = a.w * a.h + b.w * b.h - i;
  return u;
}

static float BoxIOU(Box a, Box b) {
  return BoxIntersection(a, b) / BoxUnion(a, b);
}

static void PostNMS(std::vector<Detection>* det_bboxes,
                    float thresh,
                    int classes) {
  PADDLE_ENFORCE_LE_INT_MAX(det_bboxes->size(), "detection boxes size");
  int total = static_cast<int>(det_bboxes->size());
  if (total <= 0) {
    return;
  }

  Detection* dets = det_bboxes->data();
  int i, j, k;
  k = total - 1;
  for (i = 0; i <= k; ++i) {
    if (dets[i].objectness == 0) {
      Detection swap = dets[i];
      dets[i] = dets[k];
      dets[k] = swap;
      --k;
      --i;
    }
  }
  total = k + 1;

  qsort(dets, total, sizeof(Detection), NMSComparator);

  for (i = 0; i < total; ++i) {
    if (dets[i].objectness == 0) continue;
    Box a = dets[i].bbox;
    for (j = i + 1; j < total; ++j) {
      if (dets[j].objectness == 0) continue;
      if (dets[j].max_prob_class_index != dets[i].max_prob_class_index) break;
      Box b = dets[j].bbox;
      if (BoxIOU(a, b) > thresh) {
        dets[j].objectness = 0;
        for (k = 0; k < classes; ++k) {
          dets[j].prob[k] = 0;
        }
      }
    }
  }
}

__global__ void YoloBoxNum(const float* input,
                           int* bbox_count,
                           const int grid_size,
                           const int class_num,
                           const int anchors_num,
                           float prob_thresh) {
  int64_t x_id =
      static_cast<int64_t>(blockIdx.x) * static_cast<int64_t>(blockDim.x) +
      static_cast<int64_t>(threadIdx.x);
  int64_t y_id =
      static_cast<int64_t>(blockIdx.y) * static_cast<int64_t>(blockDim.y) +
      static_cast<int64_t>(threadIdx.y);
  int64_t z_id =
      static_cast<int64_t>(blockIdx.z) * static_cast<int64_t>(blockDim.z) +
      static_cast<int64_t>(threadIdx.z);
  if ((x_id >= grid_size) || (y_id >= grid_size) || (z_id >= anchors_num)) {
    return;
  }

  const int grids_num = grid_size * grid_size;
  const int bbindex = y_id * grid_size + x_id;
  float objectness = input[bbindex + grids_num * (z_id * (5 + class_num) + 4)];
  if (objectness < prob_thresh) {
    return;
  }

  atomicAdd(bbox_count, 1);
}

__global__ void YoloTensorParseKernel(const float* input,
                                      const float* im_shape_data,
                                      const float* im_scale_data,
                                      float* output,
                                      int* bbox_index,
                                      const int grid_size,
                                      const int class_num,
                                      const int anchors_num,
                                      const int netw,
                                      const int neth,
                                      int* biases,
                                      float prob_thresh) {
  int64_t x_id =
      static_cast<int64_t>(blockIdx.x) * static_cast<int64_t>(blockDim.x) +
      static_cast<int64_t>(threadIdx.x);
  int64_t y_id =
      static_cast<int64_t>(blockIdx.y) * static_cast<int64_t>(blockDim.y) +
      static_cast<int64_t>(threadIdx.y);
  int64_t z_id =
      static_cast<int64_t>(blockIdx.z) * static_cast<int64_t>(blockDim.z) +
      static_cast<int64_t>(threadIdx.z);
  if ((x_id >= grid_size) || (y_id >= grid_size) || (z_id >= anchors_num)) {
    return;
  }

  const float pic_h = im_shape_data[0] / im_scale_data[0];
  const float pic_w = im_shape_data[1] / im_scale_data[1];
  const int grids_num = grid_size * grid_size;
  const int bbindex = y_id * grid_size + x_id;
  float objectness = input[bbindex + grids_num * (z_id * (5 + class_num) + 4)];
  if (objectness < prob_thresh) {
    return;
  }

  int cur_bbox_index = atomicAdd(bbox_index, 1);
  int tensor_index = cur_bbox_index * (5 + class_num);

  // x
  float x = input[bbindex + grids_num * (z_id * (5 + class_num) + 0)];
  x = (x + static_cast<float>(x_id)) * static_cast<float>(pic_w) /
      static_cast<float>(grid_size);
  // y
  float y = input[bbindex + grids_num * (z_id * (5 + class_num) + 1)];
  y = (y + static_cast<float>(y_id)) * static_cast<float>(pic_h) /
      static_cast<float>(grid_size);
  // w
  float w = input[bbindex + grids_num * (z_id * (5 + class_num) + 2)];
  w = w * biases[2 * z_id] * pic_w / netw;
  // h
  float h = input[bbindex + grids_num * (z_id * (5 + class_num) + 3)];
  h = h * biases[2 * z_id + 1] * pic_h / neth;

  output[tensor_index] = objectness;
  output[tensor_index + 1] = x - w / 2;
  output[tensor_index + 2] = y - h / 2;
  output[tensor_index + 3] = x + w / 2;
  output[tensor_index + 4] = y + h / 2;
  output[tensor_index + 1] =
      output[tensor_index + 1] > 0 ? output[tensor_index + 1] : 0.f;
  output[tensor_index + 2] =
      output[tensor_index + 2] > 0 ? output[tensor_index + 2] : 0.f;
  output[tensor_index + 3] = output[tensor_index + 3] < pic_w - 1
                                 ? output[tensor_index + 3]
                                 : pic_w - 1;
  output[tensor_index + 4] = output[tensor_index + 4] < pic_h - 1
                                 ? output[tensor_index + 4]
                                 : pic_h - 1;

  // Probabilities of classes
  for (int i = 0; i < class_num; ++i) {
    float prob =
        input[bbindex + grids_num * (z_id * (5 + class_num) + (5 + i))] *
        objectness;
    output[tensor_index + 5 + i] = prob;
  }
}

static void YoloTensorParseCuda(
    const float* input_data,  // [in] YOLO_BOX_HEAD layer output
    const float* image_shape_data,
    const float* image_scale_data,
    float** bboxes_tensor_ptr,  // [out] Bounding boxes output tensor
    int* bbox_count_max_alloc,  // [in/out] maximum bounding Box number
                                // allocated in dev
    int* bbox_count_host,  // [in/out] bounding boxes number recorded in host
    int* bbox_count_device_ptr,  // [in/out] bounding boxes number calculated
                                 // in
                                 // device side
    int* bbox_index_device_ptr,  // [in] bounding Box index for kernel threads
                                 // shared access
    int grid_size,
    int class_num,
    int anchors_num,
    int netw,
    int neth,
    int* biases_device,
    float prob_thresh) {
  dim3 threads_per_block(16, 16, 4);
  dim3 number_of_blocks((grid_size / threads_per_block.x) + 1,
                        (grid_size / threads_per_block.y) + 1,
                        (anchors_num / threads_per_block.z) + 1);

  // Estimate how many boxes will be chosen
  int bbox_count = 0;
#ifdef PADDLE_WITH_HIP
  hipMemcpy(
      bbox_count_device_ptr, &bbox_count, sizeof(int), hipMemcpyHostToDevice);
#else
  cudaMemcpy(
      bbox_count_device_ptr, &bbox_count, sizeof(int), cudaMemcpyHostToDevice);
#endif
  YoloBoxNum<<<number_of_blocks, threads_per_block, 0>>>(input_data,
                                                         bbox_count_device_ptr,
                                                         grid_size,
                                                         class_num,
                                                         anchors_num,
                                                         prob_thresh);
#ifdef PADDLE_WITH_HIP
  hipMemcpy(
      &bbox_count, bbox_count_device_ptr, sizeof(int), hipMemcpyDeviceToHost);
#else
  cudaMemcpy(
      &bbox_count, bbox_count_device_ptr, sizeof(int), cudaMemcpyDeviceToHost);
#endif

  // Record actual bbox number
  *bbox_count_host = bbox_count;

  // Obtain previous allocated bbox tensor in device side
  float* bbox_tensor = *bboxes_tensor_ptr;
  // Update previous maximum bbox number
  if (bbox_count > *bbox_count_max_alloc) {
#ifdef PADDLE_WITH_HIP
    hipFree(bbox_tensor);
    hipMalloc(&bbox_tensor, bbox_count * (5 + class_num) * sizeof(float));
#else
    cudaFree(bbox_tensor);
    cudaMalloc(&bbox_tensor, bbox_count * (5 + class_num) * sizeof(float));
#endif
    *bbox_count_max_alloc = bbox_count;
    *bboxes_tensor_ptr = bbox_tensor;
  }

  // Now generate bboxes
  int bbox_index = 0;
#ifdef PADDLE_WITH_HIP
  hipMemcpy(
      bbox_index_device_ptr, &bbox_index, sizeof(int), hipMemcpyHostToDevice);
#else
  cudaMemcpy(
      bbox_index_device_ptr, &bbox_index, sizeof(int), cudaMemcpyHostToDevice);
#endif
  YoloTensorParseKernel<<<number_of_blocks, threads_per_block, 0>>>(
      input_data,
      image_shape_data,
      image_scale_data,
      bbox_tensor,
      bbox_index_device_ptr,
      grid_size,
      class_num,
      anchors_num,
      netw,
      neth,
      biases_device,
      prob_thresh);
}

template <typename T, typename Context>
void YoloBoxPostKernel(const Context& dev_ctx,
                       const DenseTensor& boxes0,
                       const DenseTensor& boxes1,
                       const DenseTensor& boxes2,
                       const DenseTensor& image_shape,
                       const DenseTensor& image_scale,
                       const std::vector<int>& anchors0,
                       const std::vector<int>& anchors1,
                       const std::vector<int>& anchors2,
                       int class_num,
                       float conf_thresh,
                       int downsample_ratio0,
                       int downsample_ratio1,
                       int downsample_ratio2,
                       bool clip_bbox UNUSED,
                       float scale_x_y UNUSED,
                       float nms_threshold,
                       DenseTensor* out,
                       DenseTensor* nms_rois_num) {
  // prepare inputs
  std::vector<const float*> boxes_input(3);
  std::vector<std::vector<int32_t>> boxes_input_dims(3);
  const DenseTensor* boxes_tensor;
  for (int i = 0; i < 3; i++) {
    if (i == 0) {
      boxes_tensor = &boxes0;
    } else if (i == 1) {
      boxes_tensor = &boxes1;
    } else if (i == 2) {
      boxes_tensor = &boxes2;
    }
    boxes_input[i] = boxes_tensor->data<float>();
    auto dims = boxes_tensor->dims();
    for (int j = 0; j < dims.size(); j++) {
      boxes_input_dims[i].push_back(dims[j]);
    }
  }
  const float* image_shape_data = image_shape.data<float>();
  const float* image_scale_data = image_scale.data<float>();

  // prepare outputs
  auto* boxes_scores_tensor = out;
  auto* boxes_num_tensor = nms_rois_num;

  // prepare anchors
  std::vector<int32_t> anchors;
  anchors.insert(anchors.end(), anchors0.begin(), anchors0.end());
  anchors.insert(anchors.end(), anchors1.begin(), anchors1.end());
  anchors.insert(anchors.end(), anchors2.begin(), anchors2.end());
  int* device_anchors;
#ifdef PADDLE_WITH_HIP
  hipMalloc(reinterpret_cast<void**>(&device_anchors),
            anchors.size() * sizeof(int));
  hipMemcpy(device_anchors,
            anchors.data(),
            anchors.size() * sizeof(int),
            hipMemcpyHostToDevice);
#else
  cudaMalloc(reinterpret_cast<void**>(&device_anchors),
             anchors.size() * sizeof(int));
  cudaMemcpy(device_anchors,
             anchors.data(),
             anchors.size() * sizeof(int),
             cudaMemcpyHostToDevice);
#endif
  int* device_anchors_ptr[3];
  device_anchors_ptr[0] = device_anchors;
  device_anchors_ptr[1] = device_anchors_ptr[0] + anchors0.size();
  device_anchors_ptr[2] = device_anchors_ptr[1] + anchors1.size();
  std::vector<int> anchors_num{static_cast<int>(anchors0.size()) / 2,
                               static_cast<int>(anchors1.size()) / 2,
                               static_cast<int>(anchors2.size()) / 2};

  // prepare other attrs
  std::vector<int> downsample_ratio{
      downsample_ratio0, downsample_ratio1, downsample_ratio2};
  // clip_bbox and scale_x_y is not used now!

  int64_t batch = image_shape.dims()[0];

  TensorInfo* ts_info = new TensorInfo[batch * boxes_input.size()];
  for (int64_t i = 0; i < batch * boxes_input.size(); i++) {
#ifdef PADDLE_WITH_HIP
    hipMalloc(
        reinterpret_cast<void**>(&ts_info[i].bboxes_dev_ptr),
        ts_info[i].bbox_count_max_alloc * (5 + class_num) * sizeof(float));
#else
    cudaMalloc(
        reinterpret_cast<void**>(&ts_info[i].bboxes_dev_ptr),
        ts_info[i].bbox_count_max_alloc * (5 + class_num) * sizeof(float));
#endif
    ts_info[i].bboxes_host_ptr = reinterpret_cast<float*>(malloc(
        ts_info[i].bbox_count_max_alloc * (5 + class_num) * sizeof(float)));
#ifdef PADDLE_WITH_HIP
    hipMalloc(reinterpret_cast<void**>(&ts_info[i].bbox_count_device_ptr),
              sizeof(int));
#else
    cudaMalloc(reinterpret_cast<void**>(&ts_info[i].bbox_count_device_ptr),
               sizeof(int));
#endif
  }

  // Box index counter in gpu memory
  // *bbox_index_device_ptr used by atomicAdd
  int* bbox_index_device_ptr;
#ifdef PADDLE_WITH_HIP
  hipMalloc(reinterpret_cast<void**>(&bbox_index_device_ptr), sizeof(int));
#else
  cudaMalloc(reinterpret_cast<void**>(&bbox_index_device_ptr), sizeof(int));
#endif

  int total_bbox = 0;
  for (int batch_id = 0; batch_id < batch; batch_id++) {
    for (int input_id = 0; input_id < boxes_input.size(); input_id++) {
      int c = boxes_input_dims[input_id][1];
      int h = boxes_input_dims[input_id][2];
      int w = boxes_input_dims[input_id][3];
      int ts_id = batch_id * boxes_input.size() + input_id;
      int bbox_count_max_alloc = ts_info[ts_id].bbox_count_max_alloc;

      YoloTensorParseCuda(
          boxes_input[input_id] + batch_id * c * h * w,
          image_shape_data + batch_id * 2,
          image_scale_data + batch_id * 2,
          &(ts_info[ts_id].bboxes_dev_ptr),   // output in gpu,must use 2-level
                                              // pointer, because we may
                                              // re-malloc
          &bbox_count_max_alloc,              // bbox_count_alloc_ptr boxes we
                                              // pre-allocate
          &(ts_info[ts_id].bbox_count_host),  // record bbox numbers
          ts_info[ts_id].bbox_count_device_ptr,  // for atomicAdd
          bbox_index_device_ptr,                 // for atomicAdd
          h,
          class_num,
          anchors_num[input_id],
          downsample_ratio[input_id] * h,
          downsample_ratio[input_id] * w,
          device_anchors_ptr[input_id],
          conf_thresh);

      // batch info update
      if (bbox_count_max_alloc > ts_info[ts_id].bbox_count_max_alloc) {
        ts_info[ts_id].bbox_count_max_alloc = bbox_count_max_alloc;
        ts_info[ts_id].bboxes_host_ptr = reinterpret_cast<float*>(
            realloc(ts_info[ts_id].bboxes_host_ptr,
                    bbox_count_max_alloc * (5 + class_num) * sizeof(float)));
      }
// we need copy bbox_count_host boxes to cpu memory
#ifdef PADDLE_WITH_HIP
      hipMemcpyAsync(
          ts_info[ts_id].bboxes_host_ptr,
          ts_info[ts_id].bboxes_dev_ptr,
          ts_info[ts_id].bbox_count_host * (5 + class_num) * sizeof(float),
          hipMemcpyDeviceToHost);
#else
      cudaMemcpyAsync(
          ts_info[ts_id].bboxes_host_ptr,
          ts_info[ts_id].bboxes_dev_ptr,
          ts_info[ts_id].bbox_count_host * (5 + class_num) * sizeof(float),
          cudaMemcpyDeviceToHost);
#endif
      total_bbox += ts_info[ts_id].bbox_count_host;
    }
  }

  boxes_scores_tensor->Resize({total_bbox > 0 ? total_bbox : 1, 6});
  float* boxes_scores_data =
      dev_ctx.template HostAlloc<float>(boxes_scores_tensor);
  memset(boxes_scores_data, 0, sizeof(float) * 6);
  boxes_num_tensor->Resize({batch});
  int* boxes_num_data = dev_ctx.template HostAlloc<int>(boxes_num_tensor);
  int boxes_scores_id = 0;

  // NMS
  for (int batch_id = 0; batch_id < batch; batch_id++) {
    std::vector<Detection> bbox_det_vec;
    for (int input_id = 0; input_id < boxes_input.size(); input_id++) {
      int ts_id = batch_id * boxes_input.size() + input_id;
      int bbox_count = ts_info[ts_id].bbox_count_host;
      if (bbox_count <= 0) {
        continue;
      }

      float* bbox_host_ptr = ts_info[ts_id].bboxes_host_ptr;
      for (int bbox_index = 0; bbox_index < bbox_count; ++bbox_index) {
        Detection bbox_det;
        memset(&bbox_det, 0, sizeof(Detection));
        bbox_det.objectness = bbox_host_ptr[bbox_index * (5 + class_num) + 0];
        bbox_det.bbox.x = bbox_host_ptr[bbox_index * (5 + class_num) + 1];
        bbox_det.bbox.y = bbox_host_ptr[bbox_index * (5 + class_num) + 2];
        bbox_det.bbox.w =
            bbox_host_ptr[bbox_index * (5 + class_num) + 3] - bbox_det.bbox.x;
        bbox_det.bbox.h =
            bbox_host_ptr[bbox_index * (5 + class_num) + 4] - bbox_det.bbox.y;
        bbox_det.classes = class_num;
        bbox_det.prob =
            reinterpret_cast<float*>(malloc(class_num * sizeof(float)));
        int max_prob_class_id = -1;
        float max_class_prob = 0.0;
        for (int class_id = 0; class_id < class_num; class_id++) {
          float prob =
              bbox_host_ptr[bbox_index * (5 + class_num) + 5 + class_id];
          bbox_det.prob[class_id] = prob;
          if (prob > max_class_prob) {
            max_class_prob = prob;
            max_prob_class_id = class_id;
          }
        }
        bbox_det.max_prob_class_index = max_prob_class_id;
        bbox_det.sort_class = max_prob_class_id;
        bbox_det_vec.push_back(bbox_det);
      }
    }
    PostNMS(&bbox_det_vec, nms_threshold, class_num);
    PADDLE_ENFORCE_LE_INT_MAX(bbox_det_vec.size(), "bbox_det_num");
    const int bbox_det_num = static_cast<int>(bbox_det_vec.size());
    for (int i = 0; i < bbox_det_num; i++) {
      boxes_scores_data[boxes_scores_id++] =
          bbox_det_vec[i].max_prob_class_index;
      boxes_scores_data[boxes_scores_id++] = bbox_det_vec[i].objectness;
      boxes_scores_data[boxes_scores_id++] = bbox_det_vec[i].bbox.x;
      boxes_scores_data[boxes_scores_id++] = bbox_det_vec[i].bbox.y;
      boxes_scores_data[boxes_scores_id++] =
          bbox_det_vec[i].bbox.w + bbox_det_vec[i].bbox.x;
      boxes_scores_data[boxes_scores_id++] =
          bbox_det_vec[i].bbox.h + bbox_det_vec[i].bbox.y;
      free(bbox_det_vec[i].prob);
    }
    boxes_num_data[batch_id] = bbox_det_num;
  }

#ifdef PADDLE_WITH_HIP
  hipFree(bbox_index_device_ptr);
#else
  cudaFree(bbox_index_device_ptr);
#endif
  for (int i = 0; i < batch * boxes_input.size(); i++) {
#ifdef PADDLE_WITH_HIP
    hipFree(ts_info[i].bboxes_dev_ptr);
    hipFree(ts_info[i].bbox_count_device_ptr);
#else
    cudaFree(ts_info[i].bboxes_dev_ptr);
    cudaFree(ts_info[i].bbox_count_device_ptr);
#endif
    free(ts_info[i].bboxes_host_ptr);
  }
  delete[] ts_info;
}
}  // namespace phi

PD_REGISTER_KERNEL(
    yolo_box_post, GPU, ALL_LAYOUT, phi::YoloBoxPostKernel, float) {
  kernel->OutputAt(0).SetDataType(phi::DataType::FLOAT32);
  kernel->OutputAt(1).SetDataType(phi::DataType::INT32);
}