#ifndef _YOLOV5_CLERAITY_INFERENCE_H_
#define _YOLOV5_CLERAITY_INFERENCE_H_

#include <string>

#include "logging.h"
#include <opencv2/opencv.hpp>
#include <NvCaffeParser.h>
#include <NvInfer.h>
#include <NvInferPlugin.h>
#include <NvOnnxParser.h>

#include <cuda.h>
#include <cuda_runtime.h>
#include <cuda_runtime_api.h>

typedef struct _YoloV5ModelParam
{
    unsigned int uiClassNum;
    unsigned int uiClearNum;
    unsigned int uiDetSize;
    float        fScoreThreshold;
    float        fNmsThreshold;
} YoloV5ModelParam;

typedef struct _ModelCommonInfo
{
    unsigned int uiModelWidth;
    unsigned int uiModelHeight;
    unsigned int uiInputSize;
    unsigned int uiOutputSize;
    unsigned int uiChannel;
    unsigned int uiBatchSize;
    std::string  strInputBlobName;
    std::string  strOutputBlobName;
} ModelCommonInfo;

typedef struct _YoloV5ModelInfo
{
    YoloV5ModelParam yolov5ModelParam;
    ModelCommonInfo          modelCommonInfo;
} YoloV5ModelInfo;

#define LOCATIONS 4
struct alignas(float) stDetection { //alignas(float)
    float bbox[LOCATIONS];
    float class_conf;
    float clear_conf;
    int class_id;
    int clear_id;
};

class YoloV5Inference
{
public:
    YoloV5Inference();
    ~YoloV5Inference();

    int YoloV5InferenceInit(YoloV5ModelInfo &stYoloV5ModelInfo, std::string strModelName, std::string strEngineName);
    int YoloV5InferenceDeinit(void);
    int YoloV5InferenceModel(cv::Mat& frame, std::vector<stDetection>& res);
private:
    float GetResizeRatio(unsigned int img_width, unsigned int img_height, unsigned int model_width, unsigned int model_height);
    void doInference(nvinfer1::IExecutionContext& context, cudaStream_t& stream, void** buffers, unsigned int inputIndex, float* input, int inputSize,
                     unsigned int ouputIndex, float* output, int outputSize, int batchSize);
    void doInferenceV2(nvinfer1::IExecutionContext& context, cudaStream_t& stream, void** buffers, unsigned int outputIndex, float* output, int outputSize, int batchSize);
    void decode_opencv_nms(std::vector<stDetection>& res, float *output, int outSize, unsigned int detSize, unsigned int classNum, unsigned int clearNum, float confThresh, float nmsThresh = 0.5);
    void CenterResetLocation(float fResizeRatio, unsigned int orig_width, unsigned int orig_height, unsigned int input_w, unsigned int input_h, stDetection &detection);
    void xywh2xyxy(float *xywh, float * xyxy);
    cv::Mat preprocess_img(cv::Mat& img, int input_w, int input_h);
private:
    YoloV5ModelInfo stYoloV5ModelInfo_;
    cudaStream_t*      pImagePreprocessStream_ = nullptr;  // 图像预处理CUDA流
    cudaStream_t*      pInferenceModelStream_  = nullptr;  // 模型推理CUDA流

    float*             pfBuffers_[2];
    float*             pfInputData_  = nullptr;
    float*             pfOutputData_ = nullptr;

    uint8_t*           pu8ImgHost_   = nullptr;  // 相关内存分配
    uint8_t*           pu8ImgDevice_ = nullptr;

    unsigned int       uiInputIndex_ = 0, uiOutputIndex_ = 0;

    Logger*            pGLogger_               = nullptr;
    nvinfer1::IRuntime*          pRuntime_               = nullptr;
    nvinfer1::ICudaEngine*       pEngine_                = nullptr;
    nvinfer1::IExecutionContext* pContext_               = nullptr;
};

#endif  // END OF _YOLOV5_CLERAITY_INFERENCE_H_