#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_LeakyReLULayers[m_LeakyReLUCount]
//			= unique_ptr_INvPlugin(nvinfer1::plugin::createPReLUPlugin(serialData, serialLength));*/
//        ++m_LeakyReLUCount;
//        return m_LeakyReLULayers[m_LeakyReLUCount - 1].get();
//    }
//    else if (std::string(layerName).find("reorg") != std::string::npos)
//    {
//        assert(m_ReorgLayer == nullptr);
//        /*m_ReorgLayer = unique_ptr_INvPlugin(
//            nvinfer1::plugin::createYOLOReorgPlugin(serialData, serialLength));*/
//        return m_ReorgLayer.get();
//    }
//    else if (std::string(layerName).find("region") != std::string::npos)
//    {
//        assert(m_RegionLayer == nullptr);
//		/*m_RegionLayer = unique_ptr_INvPlugin(
//			 nvinfer1::plugin::createYOLORegionPlugin(serialData, serialLength));*/
//        return m_RegionLayer.get();
//    }
//    else if (std::string(layerName).find("yolo") != std::string::npos)
//    {
//        assert(m_YoloLayerCount >= 0 && m_YoloLayerCount < m_MaxYoloLayers);
//        assert(m_YoloLayers[m_YoloLayerCount] == nullptr);
//        m_YoloLayers[m_YoloLayerCount]
//            = unique_ptr_IPlugin(new YoloLayerV3(serialData, serialLength));
//        ++m_YoloLayerCount;
//        return m_YoloLayers[m_YoloLayerCount - 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_MaxYoloLayers; ++i)
//    {
//        m_YoloLayers[i].reset();
//    }
//
//    m_LeakyReLUCount = 0;
//    m_YoloLayerCount = 0;
//}

/******* Yolo Layer V3 *******/
/*****************************/
namespace nvinfer1
{
	YoloLayer::YoloLayer()
	{}

	YoloLayer::YoloLayer(const void* data, size_t length)
	{
		const char *d = static_cast<const char*>(data), *a = d;
		re(d, m_NumBoxes);
		re(d, m_NumClasses);
		re(d, _n_grid_h);
		re(d, _n_grid_w);
		re(d, m_OutputSize);
		assert(d = a + length);
	}
	void YoloLayer::serialize(void* buffer)const noexcept
	{
		char *d = static_cast<char*>(buffer), *a = d;
		wr(d, m_NumBoxes);
		wr(d, m_NumClasses);
		wr(d, _n_grid_h);
		wr(d, _n_grid_w);
		wr(d, m_OutputSize);
		assert(d == a + getSerializationSize());
	}

	bool YoloLayer::supportsFormat(DataType type, PluginFormat format) const noexcept
	{
		return (type == DataType::kFLOAT && format == PluginFormat::kLINEAR);
	}

	void YoloLayer::configureWithFormat(const Dims* inputDims, int nbInputs, const Dims* outputDims, int nbOutputs, DataType type, PluginFormat format, int maxBatchSize) noexcept
	{

	}

	IPluginV2* YoloLayer::clone() const noexcept
	{
		YoloLayer *p = new YoloLayer(m_NumBoxes,m_NumClasses,_n_grid_h,_n_grid_w);
		p->setPluginNamespace(_s_plugin_namespace.c_str());
		return p;
	}

	YoloLayer::YoloLayer(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 YoloLayer::getNbOutputs() const noexcept { return 1; }

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

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

	int YoloLayer::initialize() noexcept { return 0; }

	void YoloLayer::terminate() noexcept {}

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

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

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




	PluginFieldCollection YoloLayerPluginCreator::mFC{};
	std::vector<PluginField> YoloLayerPluginCreator::mPluginAttributes;

	YoloLayerPluginCreator::YoloLayerPluginCreator()
	{
		mPluginAttributes.clear();

		mFC.nbFields = mPluginAttributes.size();
		mFC.fields = mPluginAttributes.data();
	}

	const char* YoloLayerPluginCreator::getPluginName() const noexcept
	{
		return "YOLO_TRT";
	}

	const char* YoloLayerPluginCreator::getPluginVersion() const noexcept
	{
		return "1.0";
	}

	const PluginFieldCollection* YoloLayerPluginCreator::getFieldNames()noexcept
	{
		return &mFC;
	}

	IPluginV2* YoloLayerPluginCreator::createPlugin(const char* name, const PluginFieldCollection* fc)noexcept
	{
		YoloLayer* obj = new YoloLayer();
		obj->setPluginNamespace(mNamespace.c_str());
		return obj;
	}

	IPluginV2* YoloLayerPluginCreator::deserializePlugin(const char* name, const void* serialData, size_t serialLength)noexcept
	{
		// This object will be deleted when the network is destroyed, which will
		// call MishPlugin::destroy()
		YoloLayer* obj = new YoloLayer(serialData, serialLength);
		obj->setPluginNamespace(mNamespace.c_str());
		return obj;
	}


	void YoloLayerPluginCreator::setPluginNamespace(const char* libNamespace)noexcept
	{
		mNamespace = libNamespace;
	}

	const char* YoloLayerPluginCreator::getPluginNamespace() const noexcept
	{
		return mNamespace.c_str();
	}


}