#include "plugin_factory.h"
#include "trt_utils.h"

PluginFactory::PluginFactory() : m_ReorgLayer{nullptr}, m_RegionLayer{nullptr}
{
    for (int i = 0; i < m_MaxLeakyLayers; ++i) m_LeakyReLULayers[i] = nullptr;
}

nvinfer1::IPlugin* PluginFactory::createPlugin(const char* layerName, const void* serialData,
                                               size_t serialLength)
{
    assert(isPlugin(layerName));
    if (std::string(layerName).find("leaky") != std::string::npos)
    {
        assert(m_LeakyReLUCount >= 0 && m_LeakyReLUCount <= m_MaxLeakyLayers);
        assert(m_LeakyReLULayers[m_LeakyReLUCount] == nullptr);
        ++m_LeakyReLUCount;
        return m_LeakyReLULayers[m_LeakyReLUCount - 1].get();
    }
    else if (std::string(layerName).find("reorg") != std::string::npos)
    {
        assert(m_ReorgLayer == nullptr);
        return m_ReorgLayer.get();
    }
    else if (std::string(layerName).find("region") != std::string::npos)
    {
        assert(m_RegionLayer == nullptr);
        return m_RegionLayer.get();
    }
    else if (std::string(layerName).find("yolo") != std::string::npos)
    {
        assert(m_SLayerCount >= 0 && m_SLayerCount < m_MaxSLayers);
        assert(m_SLayers[m_SLayerCount] == nullptr);
        m_SLayers[m_SLayerCount]
            = unique_ptr_IPlugin(new SLayerV3(serialData, serialLength));
        ++m_SLayerCount;
        return m_SLayers[m_SLayerCount - 1].get();
    }
    else
    {
        std::cerr << "ERROR: Unrecognised layer : " << layerName << std::endl;
        assert(0);
        return nullptr;
    }
}

bool PluginFactory::isPlugin(const char* name)
{
    return ((std::string(name).find("leaky") != std::string::npos)
            || (std::string(name).find("reorg") != std::string::npos)
            || (std::string(name).find("region") != std::string::npos)
            || (std::string(name).find("yolo") != std::string::npos));
}

void PluginFactory::destroy()
{
    m_ReorgLayer.reset();
    m_RegionLayer.reset();

    for (int i = 0; i < m_MaxLeakyLayers; ++i)
    {
        m_LeakyReLULayers[i].reset();
    }

    for (int i = 0; i < m_MaxSLayers; ++i)
    {
        m_SLayers[i].reset();
    }

    m_LeakyReLUCount = 0;
    m_SLayerCount = 0;
}

/******* Layer V3 *******/
/*****************************/
SLayerV3::SLayerV3(const void* data, size_t length)
{
    const char *d = static_cast<const char*>(data), *a = d;
    read(d, m_NumBoxes);
    read(d, m_NumClasses);
    read(d,_n_grid_h);
    read(d,_n_grid_w);
    read(d, m_OutputSize);
    assert(d = a + length);
}

SLayerV3::SLayerV3(const uint32_t& numBoxes, const uint32_t& numClasses, const uint32_t& grid_h_,const uint32_t &grid_w_):
    m_NumBoxes(numBoxes),
    m_NumClasses(numClasses),
    _n_grid_h(grid_h_),
    _n_grid_w(grid_w_)
{
    assert(m_NumBoxes > 0);
    assert(m_NumClasses > 0);
    assert(_n_grid_h > 0);
    assert(_n_grid_w > 0);
    m_OutputSize = _n_grid_h * _n_grid_w * (m_NumBoxes * (4 + 1 + m_NumClasses));
}

int SLayerV3::getNbOutputs() const { return 1; }

nvinfer1::Dims SLayerV3::getOutputDimensions(int index, const nvinfer1::Dims* inputs,
                                                int nbInputDims)
{
    assert(index == 0);
    assert(nbInputDims == 1);
    return inputs[0];
}

void SLayerV3::configure(const nvinfer1::Dims* inputDims, int nbInputs,
                            const nvinfer1::Dims* outputDims, int nbOutputs, int maxBatchSize)
{
    assert(nbInputs == 1);
    assert(inputDims != nullptr);
}

int SLayerV3::initialize() { return 0; }

void SLayerV3::terminate() {}

size_t SLayerV3::getWorkspaceSize(int maxBatchSize) const { return 0; }

int SLayerV3::enqueue(int batchSize, const void* const* inputs, void** outputs, void* workspace,
                         cudaStream_t stream)
{
    NV_CUDA_CHECK(cudaSLayerV3(inputs[0], outputs[0], batchSize,_n_grid_h,_n_grid_w, m_NumClasses,
                                  m_NumBoxes, m_OutputSize, stream));
    return 0;
}

size_t SLayerV3::getSerializationSize()
{
    return sizeof(m_NumBoxes) + sizeof(m_NumClasses) + sizeof(_n_grid_w)+sizeof(_n_grid_h) + sizeof(m_OutputSize);
}

void SLayerV3::serialize(void* buffer)
{
    char *d = static_cast<char*>(buffer), *a = d;
    write(d, m_NumBoxes);
    write(d, m_NumClasses);
    write(d,_n_grid_h);
    write(d,_n_grid_w);
    write(d, m_OutputSize);
    assert(d == a + getSerializationSize());
}