#pragma once
#include <string>
#include <vector>
#include "NvInfer.h"
#include "../include/macros.h"
namespace nvinfer1 {
class API YoloLayerPlugin : public IPluginV2IOExt {
   public:
    YoloLayerPlugin(int classCount, int numberofpoints, float confthreshkeypoints, int netWidth, int netHeight,
                    int maxOut, bool is_segmentation, bool is_pose, const int* strides, int stridesLength);

    YoloLayerPlugin(const void* data, size_t length);
    ~YoloLayerPlugin();

    int getNbOutputs() const TRT_NOEXCEPT override { return 1; }

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

    int initialize() TRT_NOEXCEPT override;

    virtual void terminate() TRT_NOEXCEPT override {}

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

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

    virtual size_t getSerializationSize() const TRT_NOEXCEPT override;

    virtual void serialize(void* buffer) const TRT_NOEXCEPT override;

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

    const char* getPluginType() const TRT_NOEXCEPT override;

    const char* getPluginVersion() const TRT_NOEXCEPT override;

    void destroy() TRT_NOEXCEPT override;

    IPluginV2IOExt* clone() const TRT_NOEXCEPT override;

    void setPluginNamespace(const char* pluginNamespace) TRT_NOEXCEPT override;

    const char* getPluginNamespace() const TRT_NOEXCEPT override;

    nvinfer1::DataType getOutputDataType(int32_t index, nvinfer1::DataType const* inputTypes,
                                         int32_t nbInputs) const TRT_NOEXCEPT;

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

    bool canBroadcastInputAcrossBatch(int inputIndex) const TRT_NOEXCEPT override;

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

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

    void detachFromContext() TRT_NOEXCEPT override;

   private:
    void forwardGpu(const float* const* inputs, float* output, cudaStream_t stream, int mYoloV8netHeight,
                    int mYoloV8NetWidth, int batchSize);
    int mThreadCount = 256;
    const char* mPluginNamespace;
    int mClassCount;
    int mNumberofpoints;
    float mConfthreshkeypoints;
    int mYoloV8NetWidth;
    int mYoloV8netHeight;
    int mMaxOutObject;
    bool is_segmentation_;
    bool is_pose_;
    int* mStrides;
    int mStridesLength;
};

class API YoloPluginCreator : public IPluginCreator {
   public:
    YoloPluginCreator();
    ~YoloPluginCreator() override = default;

    const char* getPluginName() const TRT_NOEXCEPT override;

    const char* getPluginVersion() const TRT_NOEXCEPT override;

    const nvinfer1::PluginFieldCollection* getFieldNames() TRT_NOEXCEPT override;

    nvinfer1::IPluginV2IOExt* createPlugin(const char* name,
                                           const nvinfer1::PluginFieldCollection* fc) TRT_NOEXCEPT override;

    nvinfer1::IPluginV2IOExt* deserializePlugin(const char* name, const void* serialData,
                                                size_t serialLength) TRT_NOEXCEPT override;

    void setPluginNamespace(const char* libNamespace) TRT_NOEXCEPT override { mNamespace = libNamespace; }

    const char* getPluginNamespace() const TRT_NOEXCEPT override { return mNamespace.c_str(); }

   private:
    std::string mNamespace;
    static PluginFieldCollection mFC;
    static std::vector<PluginField> mPluginAttributes;
};
REGISTER_TENSORRT_PLUGIN(YoloPluginCreator);
}  // namespace nvinfer1