#ifndef _YOLO_LAYER_H
#define _YOLO_LAYER_H

#include <assert.h>
#include <cmath>
#include <string.h>
#include <cublas_v2.h>
#include "NvInfer.h"
#include "Utils.h"
#include <iostream>
#include "NvInferPlugin.h"

struct YoloKernel
{
    int width;
    int height;
    int everyYoloAnchors;
    float anchors[10];   // 一组yolo输出层中 anchors的数据个数 等于 3*2， 可以设置的更大一点，这个无所谓
};

struct alignas(float) Detection{
    //x y w h
    float bbox[4];
    float det_confidence;
    float class_id;
    float class_confidence;
};

namespace nvinfer1
{
    class YoloLayerPlugin: public IPluginV2IOExt
    {
        public:
            YoloLayerPlugin(const PluginFieldCollection& fc);
            YoloLayerPlugin(const void* data, size_t length);

            ~YoloLayerPlugin();

            int getNbOutputs() const override
            {
                return 1;
            }

            Dims getOutputDimensions(int index, const Dims* inputs, int nbInputDims) override;

            int initialize() override;

            virtual void terminate() override {};

            virtual size_t getWorkspaceSize(int maxBatchSize) const override { return 0;}

            virtual int enqueue(int batchSize, const void*const * inputs, void** outputs, void* workspace, cudaStream_t stream) override;

            virtual size_t getSerializationSize() const override;

            virtual void serialize(void* buffer) const override;

            bool supportsFormatCombination(int pos, const PluginTensorDesc* inOut, int nbInputs, int nbOutputs) const override {
                return inOut[pos].format == TensorFormat::kLINEAR && inOut[pos].type == DataType::kFLOAT;
            }

            const char* getPluginType() const override;

            const char* getPluginVersion() const override;

            void destroy() override;

            IPluginV2IOExt* clone() const override;

            void setPluginNamespace(const char* pluginNamespace) override;

            const char* getPluginNamespace() const override;

            DataType getOutputDataType(int index, const nvinfer1::DataType* inputTypes, int nbInputs) const override;

            bool isOutputBroadcastAcrossBatch(int outputIndex, const bool* inputIsBroadcasted, int nbInputs) const override;

            bool canBroadcastInputAcrossBatch(int inputIndex) const override;

            void attachToContext(
                    cudnnContext* cudnnContext, cublasContext* cublasContext, IGpuAllocator* gpuAllocator) override;

            void configurePlugin(const PluginTensorDesc* in, int nbInput, const PluginTensorDesc* out, int nbOutput) override;

            void detachFromContext() override;

        private:
            void forwardGpu(const float *const * inputs,float * output, cudaStream_t stream,int batchSize = 1);
            int mClassCount;        // 检测的目标的类别，从cfg文件获取，在cfg 设置
            int mInput_w;           // 图像输入的尺寸，从cfg获取
            int mInput_h;           // 由于umsample层的原因，宽度和高度要想等，TODO 调整
            int mNumYoloLayers;     // yolo输出层的数量，从cfg获取，无需设置
            std::vector<YoloKernel> mYoloKernel;

            float mIgnore_thresh = 0.4;     // 置信度阈值，可以调整
            int max_output_box = 1000;      // 最大输出数量
            int mThreadCount = 256;         // cuda 内核函数，每一block中线程数量
            const char* mPluginNamespace;   // 该插件名称

    };
    // 继承与IPluginCreator，重写虚函数
    class YoloPluginCreator : public IPluginCreator
    {
        public:
            YoloPluginCreator();

            ~YoloPluginCreator() override = default;
            const char* getPluginName() const override;
            const char* getPluginVersion() const override;
            const PluginFieldCollection* getFieldNames() override;
            // 生成插件，这个是在 build network时调用
            IPluginV2IOExt* createPlugin(const char* name, const PluginFieldCollection* fc) override;
            // 反序列化，在读取保存的trt模型engine时调用，负责解析插件
            IPluginV2IOExt* deserializePlugin(const char* name, const void* serialData, size_t serialLength) override;
            void setPluginNamespace(const char* libNamespace) override{
                mNamespace = libNamespace;
            }
            const char* getPluginNamespace() const override{
                return mNamespace.c_str();
            }
        private:
            std::string mNamespace;
    };
};

#endif