/*
 * Copyright 2018-2019 Autoware Foundation. All rights reserved.
 * Copyright (c) 2019-2021 Intel Corporation (oneAPI modifications)
 *
 * 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 <sycl/sycl.hpp>
#include <memory>
#include <vector>
#include "pointpillars/nms.hpp"
#include "pointpillars/pointpillars_util.hpp"

namespace pointpillars {

/**
 * PointPillar's PostProcessing
 *
 * Use the output of the RegionProposalNetwork and the
 * anchors generated by the AnchorGrid to decode the
 * object position, dimension and class, filter out
 * redundant/clutter objects using NMS and sort them
 * according to likelihood. Finally convert into
 * object representation.
 */
class PostProcess {
 private:
  const float float_min_;
  const float float_max_;
  const size_t num_anchor_x_inds_;
  const size_t num_anchor_y_inds_;
  const size_t num_anchor_r_inds_;
  const size_t num_cls_;
  const float score_threshold_;
  const size_t num_threads_;
  const size_t num_box_corners_;
  const size_t num_output_box_feature_;

  std::unique_ptr<NMS> nms_ptr_;

 public:
  /**
  * @brief Constructor
  * @param[in] float_min The lowest float value
  * @param[in] float_max The maximum float value
  * @param[in] num_anchor_x_inds Number of x-indexes for anchors
  * @param[in] num_anchor_y_inds Number of y-indexes for anchors
  * @param[in] num_anchor_r_inds Number of rotation-indexes for anchors
  * @param[in] score_threshold Score threshold for filtering output
  * @param[in] num_threads Number of threads when launching kernel
  * @param[in] nms_overlap_threshold IOU threshold for NMS
  * @param[in] num_box_corners Number of box's corner
  * @param[in] num_output_box_feature Number of output box's feature
  */
  PostProcess(const float float_min, const float float_max, const size_t num_anchor_x_inds,
              const size_t num_anchor_y_inds, const size_t num_anchor_r_inds, const size_t num_cls,
              const float score_threshold, const size_t num_threads, const float nms_overlap_threshold,
              const size_t num_box_corners, const size_t num_output_box_feature);

  /**
  * @brief Postprocessing for the network output
  * @param[in] rpn_box_output Box predictions from the network output
  * @param[in] rpn_cls_output Class predictions from the network output
  * @param[in] rpn_dir_output Direction predictions from the network output
  * @param[in] dev_anchor_mask Anchor mask for filtering the network output
  * @param[in] dev_anchors_px X-coordinate values for corresponding anchors
  * @param[in] dev_anchors_py Y-coordinate values for corresponding anchors
  * @param[in] dev_anchors_pz Z-coordinate values for corresponding anchors
  * @param[in] dev_anchors_dx X-dimension values for corresponding anchors
  * @param[in] dev_anchors_dy Y-dimension values for corresponding anchors
  * @param[in] dev_anchors_dz Z-dimension values for corresponding anchors
  * @param[in] dev_anchors_ro Rotation values for corresponding anchors
  * @param[in] dev_filtered_box Filtered box predictions
  * @param[in] dev_filtered_score Filtered score predictions
  * @param[in] dev_filtered_dir Filtered direction predictions
  * @param[in] dev_box_for_nms Decoded boxes in min_x min_y max_x max_y represenation from pose and dimension
  * @param[in] dev_filter_count The number of filtered output
  * @param[out] out_detection Output bounding boxes
  * @details dev_* represents device memory allocated variables
  */
  void DoPostProcess(const float *rpn_box_output, const float *rpn_cls_output, const float *rpn_dir_output,
                     int *dev_anchor_mask, const float *dev_anchors_px, const float *dev_anchors_py,
                     const float *dev_anchors_pz, const float *dev_anchors_dx, const float *dev_anchors_dy,
                     const float *dev_anchors_dz, const float *dev_anchors_ro, float *dev_multiclass_score,
                     float *dev_filtered_box, float *dev_filtered_score, int *dev_filtered_dir,
                     int *dev_filtered_class_id, float *dev_box_for_nms, int *dev_filter_count,
                     std::vector<ObjectDetection> &detections);
};
}  // namespace pointpillars
