#ifndef __PLUGIN_LAYER_H__
#define __PLUGIN_LAYER_H__

#include <cassert>
#include <cstring>
#include <cudnn.h>
#include <iostream>
#include <memory>

#include "NvInferPlugin.h"

#define NV_CUDA_CHECK(status)                                                                      \
    {                                                                                              \
        if (status != 0)                                                                           \
        {                                                                                          \
            std::cout << "Cuda failure: " << cudaGetErrorString(status) << " in file " << __FILE__ \
                      << " at line " << __LINE__ << std::endl;                                     \
            abort();                                                                               \
        }                                                                                          \
    }

cudaError_t cudaSLayerV3(const void* input, void* output, const uint32_t& batchSize,
    const uint32_t& n_grid_h_, const uint32_t& n_grid_w_,
    const uint32_t& numOutputClasses, const uint32_t& numBBoxes,
    uint64_t outputSize, cudaStream_t stream);

class PluginFactory : public nvinfer1::IPluginFactory
{

public:
    PluginFactory();
    nvinfer1::IPlugin* createPlugin(const char* layerName, const void* serialData,
                                    size_t serialLength) override;
    bool isPlugin(const char* name);
    void destroy();

private:
    static const int m_MaxLeakyLayers = 72;
    static const int m_ReorgStride = 2;
    static constexpr float m_LeakyNegSlope = 0.1f;
    static const int m_NumBoxes = 5;
    static const int m_NumCoords = 4;
    static const int m_NumClasses = 80;
    static const int m_MaxSLayers = 3;
    int m_LeakyReLUCount = 0;
    int m_SLayerCount = 0;
    nvinfer1::plugin::RegionParameters m_RegionParameters{m_NumBoxes, m_NumCoords, m_NumClasses,
                                                          nullptr};

    struct INvPluginDeleter
    {
        void operator()(nvinfer1::plugin::INvPlugin* ptr)
        {
            if (ptr)
            {
                ptr->destroy();
            }
        }
    };
    struct IPluginDeleter
    {
        void operator()(nvinfer1::IPlugin* ptr)
        {
            if (ptr)
            {
                ptr->terminate();
            }
        }
    };
    typedef std::unique_ptr<nvinfer1::plugin::INvPlugin, INvPluginDeleter> unique_ptr_INvPlugin;
    typedef std::unique_ptr<nvinfer1::IPlugin, IPluginDeleter> unique_ptr_IPlugin;

    unique_ptr_INvPlugin m_ReorgLayer;
    unique_ptr_INvPlugin m_RegionLayer;
    unique_ptr_INvPlugin m_LeakyReLULayers[m_MaxLeakyLayers];
    unique_ptr_IPlugin m_SLayers[m_MaxSLayers];
};

class SLayerV3 : public nvinfer1::IPlugin
{
public:
    SLayerV3(const void* data, size_t length);
    SLayerV3(const uint32_t& numBoxes, const uint32_t& numClasses, const uint32_t& grid_h_,const uint32_t &grid_w_);
    int getNbOutputs() const override;
    nvinfer1::Dims getOutputDimensions(int index, const nvinfer1::Dims* inputs,
                                       int nbInputDims) override;
    void configure(const nvinfer1::Dims* inputDims, int nbInputs, const nvinfer1::Dims* outputDims,
                   int nbOutputs, int maxBatchSize) override;
    int initialize() override;
    void terminate() override;
    size_t getWorkspaceSize(int maxBatchSize) const override;
    int enqueue(int batchSize, const void* const* intputs, void** outputs, void* workspace,
                cudaStream_t stream) override;
    size_t getSerializationSize() override;
    void serialize(void* buffer) override;

private:
    template <typename T>
    void write(char*& buffer, const T& val)
    {
        *reinterpret_cast<T*>(buffer) = val;
        buffer += sizeof(T);
    }

    template <typename T>
    void read(const char*& buffer, T& val)
    {
        val = *reinterpret_cast<const T*>(buffer);
        buffer += sizeof(T);
    }
    uint32_t m_NumBoxes;
    uint32_t m_NumClasses;
    uint32_t m_GridSize;
    uint64_t m_OutputSize;
    uint32_t _n_grid_h;
    uint32_t _n_grid_w;
};

#endif // __PLUGIN_LAYER_H__