#include "model.h"
#include <memory>
#include <vector>
#include <fstream>
#include <iostream>
#include <sstream>
#include <iomanip>

using namespace nvinfer1;
// REGISTER_TENSORRT_PLUGIN(DetectPluginCreator);


Detecter::Detecter( const NetworkInfo& networkInfo, const InferParams& inferParams, int type) :
    m_NetworkType(networkInfo.networkType),
    m_InputBlobName(networkInfo.inputBlobName),
    m_InputH(416),
    m_InputW(416),
    m_InputC(3),
    m_InputSize(m_InputH*m_InputW*m_InputC),
    m_ProbThresh(inferParams.probThresh),
    m_NMSThresh(inferParams.nmsThresh),
    m_Logger(Logger()),
    m_Network(nullptr),
    m_Builder(nullptr),
    m_ModelStream(nullptr),
    m_Engine(nullptr),
    m_Context(nullptr),
    m_InputBindingIndex(-1),
    m_CudaStream(nullptr),
    m_PluginFactory(new PluginFactory),
    m_TinyMaxpoolPaddingFormula(new YoloTinyMaxpoolPaddingFormula)
{
    m_EnginePath = m_staticStruct::model_path;
    if(!fileExists(m_EnginePath))
    {
        m_configBlocks = parseConfigFile(m_staticStruct::model_cfg);
        parseConfigBlocks();
        createYOLOEngine();
    }
    setOutput(type);
    DEBUG((boost::format("m_EnginePath:%s")%m_EnginePath).str());
    assert(m_PluginFactory != nullptr);
    m_Engine = loadTRTEngine(m_EnginePath, m_PluginFactory, m_Logger);
    assert(m_Engine != nullptr);
    m_Context = m_Engine->createExecutionContext();
    assert(m_Context != nullptr);
    m_InputBindingIndex = m_Engine->getBindingIndex(m_InputBlobName.c_str());
    assert(m_InputBindingIndex != -1);
    assert(m_BatchSize <= static_cast<uint32_t>(m_Engine->getMaxBatchSize()));
    allocateBuffers();
    NV_CUDA_CHECK(cudaStreamCreate(&m_CudaStream));
    assert(verifyEngine());
}

Detecter::~Detecter()
{
    for (auto& tensor : m_OutputTensors) NV_CUDA_CHECK(cudaFreeHost(tensor.hostBuffer));
    for (auto& deviceBuffer : m_DeviceBuffers) NV_CUDA_CHECK(cudaFree(deviceBuffer));
    NV_CUDA_CHECK(cudaStreamDestroy(m_CudaStream));
    if (m_Context)
    {
        m_Context->destroy();
        m_Context = nullptr;
    }

    if (m_Engine)
    {
        m_Engine->destroy();
        m_Engine = nullptr;
    }

    if (m_PluginFactory)
    {
        m_PluginFactory->destroy();
        m_PluginFactory = nullptr;
    }

    m_TinyMaxpoolPaddingFormula.reset();
}

std::vector<std::map<std::string, std::string>> Detecter::parseConfigFile(const std::string cfgFilePath)
{
    std::cout << "::::::::::" << cfgFilePath <<std::endl;
    assert(fileExists(cfgFilePath));
    std::ifstream file(cfgFilePath);
    assert(file.good());
    std::string line;
    std::vector<std::map<std::string, std::string>> blocks;
    std::map<std::string, std::string> block;

    while (getline(file, line))
    {
        if (line.empty()) continue;
        if (line.front() == '#') continue;
        line = trim(line);
        if (line.front() == '[')
        {
            if (!block.empty())
            {
                blocks.push_back(block);
                block.clear();
            }
            std::string key = "type";
            std::string value = trim(line.substr(1, line.size() - 2));
            block.insert(std::pair<std::string, std::string>(key, value));
        }
        else
        {
            size_t cpos = line.find('=');
            std::string key = trim(line.substr(0, cpos));
            std::string value = trim(line.substr(cpos + 1));
            block.insert(std::pair<std::string, std::string>(key, value));
        }
    }
    blocks.push_back(block);
    return blocks;
}

void Detecter::parseConfigBlocks()
{
    for (auto block : m_configBlocks)
    {
        if (block.at("type") == "net")
        {
            assert((block.find("height") != block.end())
                   && "Missing 'height' param in network cfg");
            assert((block.find("width") != block.end()) && "Missing 'width' param in network cfg");
            assert((block.find("channels") != block.end())
                   && "Missing 'channels' param in network cfg");
            assert((block.find("batch") != block.end())
                   && "Missing 'batch' param in network cfg");

            m_InputH = std::stoul(trim(block.at("height")));
            m_InputW = std::stoul(trim(block.at("width")));
            m_InputC = std::stoul(trim(block.at("channels")));
            m_BatchSize = std::stoi(trim(block.at("batch")));
         //   assert(m_InputW == m_InputH);
            m_InputSize = m_InputC * m_InputH * m_InputW;
        }
        else if ((block.at("type") == "region") || (block.at("type") == "yolo"))
        {
            assert((block.find("num") != block.end())
                   && std::string("Missing 'num' param in " + block.at("type") + " layer").c_str());
            assert((block.find("classes") != block.end())
                   && std::string("Missing 'classes' param in " + block.at("type") + " layer")
                          .c_str());
            assert((block.find("anchors") != block.end())
                   && std::string("Missing 'anchors' param in " + block.at("type") + " layer")
                          .c_str());

            TensorInfo outputTensor;
            std::string anchorString = block.at("anchors");
            while (!anchorString.empty())
            {
                size_t npos = anchorString.find_first_of(',');
                if (npos != std::string::npos)
                {
                    float anchor = std::stof(trim(anchorString.substr(0, npos)));
                    outputTensor.anchors.push_back(anchor);
                    anchorString.erase(0, npos + 1);
                }
                else
                {
                    float anchor = std::stof(trim(anchorString));
                    outputTensor.anchors.push_back(anchor);
                    break;
                }
            }

            assert((block.find("mask") != block.end())
                   && std::string("Missing 'mask' param in " + block.at("type") + " layer")
                          .c_str());

            std::string maskString = block.at("mask");
            while (!maskString.empty())
            {
                size_t npos = maskString.find_first_of(',');
                if (npos != std::string::npos)
                {
                    uint32_t mask = std::stoul(trim(maskString.substr(0, npos)));
                    outputTensor.masks.push_back(mask);
                    maskString.erase(0, npos + 1);
                }
                else
                {
                    uint32_t mask = std::stoul(trim(maskString));
                    outputTensor.masks.push_back(mask);
                    break;
                }
            }

            outputTensor.numBBoxes = outputTensor.masks.size() > 0
                ? outputTensor.masks.size()
                : std::stoul(trim(block.at("num")));
            outputTensor.numClasses = std::stoul(block.at("classes"));
            if (m_ClassNames.empty())
            {
                for (uint32_t i=0;i< outputTensor.numClasses;++i)
                {
                    m_ClassNames.push_back(std::to_string(i));
                }
            }
            outputTensor.blobName = "yolo_" + std::to_string(_n_yolo_ind);
            outputTensor.gridSize = (m_InputH / 32) * pow(2, _n_yolo_ind);
            outputTensor.grid_h = (m_InputH / 32) * pow(2, _n_yolo_ind);
            outputTensor.grid_w = (m_InputW / 32) * pow(2, _n_yolo_ind);
            outputTensor.gridSize = (m_InputH / 32) * pow(2, 2-_n_yolo_ind);
            outputTensor.grid_h = (m_InputH / 32) * pow(2, 2-_n_yolo_ind);
            outputTensor.grid_w = (m_InputW / 32) * pow(2, 2-_n_yolo_ind);
            outputTensor.stride = m_InputH / outputTensor.gridSize;
            outputTensor.stride_h = m_InputH / outputTensor.grid_h;
            outputTensor.stride_w = m_InputW / outputTensor.grid_w;
            outputTensor.volume = outputTensor.grid_h* outputTensor.grid_w
                *(outputTensor.numBBoxes*(5 + outputTensor.numClasses));
            m_OutputTensors.push_back(outputTensor);
            _n_yolo_ind++;
        }
    }
}

void Detecter::createYOLOEngine(const nvinfer1::DataType dataType, Int8EntropyCalibrator* calibrator)
{
    if (fileExists(m_EnginePath))return;
    std::vector<float> weights = loadWeights(m_staticStruct::model_wts, m_NetworkType);
    std::vector<nvinfer1::Weights> trtWeights;
    int weightPtr = 0;
    int channels = m_InputC;
    m_Builder = nvinfer1::createInferBuilder(m_Logger);
    nvinfer1::IBuilderConfig* config = m_Builder->createBuilderConfig();
    m_Network = m_Builder->createNetworkV2(0U);
    if ((dataType == nvinfer1::DataType::kINT8 && !m_Builder->platformHasFastInt8())
        || (dataType == nvinfer1::DataType::kHALF && !m_Builder->platformHasFastFp16()))
    {
        std::cout << "Platform doesn't support this precision." << std::endl;
        assert(0);
    }

    nvinfer1::ITensor* data = m_Network->addInput(
        m_InputBlobName.c_str(), nvinfer1::DataType::kFLOAT,
        nvinfer1::DimsCHW{static_cast<int>(m_InputC), static_cast<int>(m_InputH),
                          static_cast<int>(m_InputW)});
    assert(data != nullptr);
    // Add elementwise layer to normalize pixel values 0-1
    nvinfer1::Dims divDims{
        3,
        {static_cast<int>(m_InputC), static_cast<int>(m_InputH), static_cast<int>(m_InputW)},
        {nvinfer1::DimensionType::kCHANNEL, nvinfer1::DimensionType::kSPATIAL,
         nvinfer1::DimensionType::kSPATIAL}};
    nvinfer1::Weights divWeights{nvinfer1::DataType::kFLOAT, nullptr,
                                 static_cast<int64_t>(m_InputSize)};
    float* divWt = new float[m_InputSize];
    for (uint32_t w = 0; w < m_InputSize; ++w) divWt[w] = 255.0;
    divWeights.values = divWt;
    trtWeights.push_back(divWeights);
    nvinfer1::IConstantLayer* constDivide = m_Network->addConstant(divDims, divWeights);
    assert(constDivide != nullptr);
    nvinfer1::IElementWiseLayer* elementDivide = m_Network->addElementWise(
        *data, *constDivide->getOutput(0), nvinfer1::ElementWiseOperation::kDIV);
    assert(elementDivide != nullptr);

    nvinfer1::ITensor* previous = elementDivide->getOutput(0);
    std::vector<nvinfer1::ITensor*> tensorOutputs;
    uint32_t outputTensorCount = 0;

    // build the network using the network API
    for (uint32_t i = 0; i < m_configBlocks.size(); ++i)
    {
        // check if num. of channels is correct
        assert(getNumChannels(previous) == channels);
        std::string layerIndex = "(" + std::to_string(i) + ")";

        if (m_configBlocks.at(i).at("type") == "net")
        {
            printLayerInfo("", "layer", "     inp_size", "     out_size", "weightPtr");
        }
        else if (m_configBlocks.at(i).at("type") == "convolutional")
        {
            std::string inputVol = dimsToString(previous->getDimensions());
            nvinfer1::ILayer* out;
            std::string layerType;
            //check activation
            std::string activation = "";
            if (m_configBlocks.at(i).find("activation") != m_configBlocks.at(i).end())
            {
                activation = m_configBlocks[i]["activation"];
            }
            // check if batch_norm enabled
            if ((m_configBlocks.at(i).find("batch_normalize") != m_configBlocks.at(i).end()) &&
                ("leaky" == activation))
            {
                out = netAddConvBNLeaky(i, m_configBlocks.at(i), weights, trtWeights, weightPtr,
                                        channels, previous, m_Network);
                layerType = "conv-bn-leaky";
            }
            else if ((m_configBlocks.at(i).find("batch_normalize") != m_configBlocks.at(i).end()) &&
                ("mish" == activation))
            {
                out = net_conv_bn_mish(i, m_configBlocks.at(i), weights, trtWeights, weightPtr,
                                        channels, previous, m_Network);
                layerType = "conv-bn-mish";
            }
            else// if("linear" == activation)
            {
                out = netAddConvLinear(i, m_configBlocks.at(i), weights, trtWeights, weightPtr,
                                       channels, previous, m_Network);
                layerType = "conv-linear";
            }
            previous = out->getOutput(0);
            assert(previous != nullptr);
            channels = getNumChannels(previous);
            std::string outputVol = dimsToString(previous->getDimensions());
            tensorOutputs.push_back(out->getOutput(0));
            printLayerInfo(layerIndex, layerType, inputVol, outputVol, std::to_string(weightPtr));
        }
        else if (m_configBlocks.at(i).at("type") == "shortcut")
        {
            assert(m_configBlocks.at(i).at("activation") == "linear");
            assert(m_configBlocks.at(i).find("from") != m_configBlocks.at(i).end());
            int from = stoi(m_configBlocks.at(i).at("from"));

            std::string inputVol = dimsToString(previous->getDimensions());
            // check if indexes are correct
            assert((i - 2 >= 0) && (i - 2 < tensorOutputs.size()));
            assert((i + from - 1 >= 0) && (i + from - 1 < tensorOutputs.size()));
            assert(i + from - 1 < i - 2);
            nvinfer1::IElementWiseLayer* ew
                = m_Network->addElementWise(*tensorOutputs[i - 2], *tensorOutputs[i + from - 1],
                                            nvinfer1::ElementWiseOperation::kSUM);
            assert(ew != nullptr);
            std::string ewLayerName = "shortcut_" + std::to_string(i);
            ew->setName(ewLayerName.c_str());
            previous = ew->getOutput(0);
            assert(previous != nullptr);
            std::string outputVol = dimsToString(previous->getDimensions());
            tensorOutputs.push_back(ew->getOutput(0));
            printLayerInfo(layerIndex, "skip", inputVol, outputVol, "    -");
        }
        else if (m_configBlocks.at(i).at("type") == "yolo")
        {
            nvinfer1::Dims prevTensorDims = previous->getDimensions();
           // assert(prevTensorDims.d[1] == prevTensorDims.d[2]);
            TensorInfo& curYoloTensor = m_OutputTensors.at(outputTensorCount);
            curYoloTensor.gridSize = prevTensorDims.d[1];
            curYoloTensor.grid_h = prevTensorDims.d[1];
            curYoloTensor.grid_w = prevTensorDims.d[2];
            curYoloTensor.stride = m_InputW / curYoloTensor.gridSize;
            curYoloTensor.stride_h = m_InputH / curYoloTensor.grid_h;
            curYoloTensor.stride_w = m_InputW / curYoloTensor.grid_w;
            m_OutputTensors.at(outputTensorCount).volume = curYoloTensor.grid_h
                * curYoloTensor.grid_w
                * (curYoloTensor.numBBoxes * (5 + curYoloTensor.numClasses));
            std::string layerName = "yolo_" + std::to_string(outputTensorCount);
            curYoloTensor.blobName = layerName;
            nvinfer1::IPlugin* yoloPlugin
                = new SLayerV3(m_OutputTensors.at(outputTensorCount).numBBoxes,
                                  m_OutputTensors.at(outputTensorCount).numClasses,
                                  m_OutputTensors.at(outputTensorCount).grid_h,
                                  m_OutputTensors.at(outputTensorCount).grid_w);
            assert(yoloPlugin != nullptr);
            nvinfer1::IPluginLayer* yolo = m_Network->addPlugin(&previous, 1, *yoloPlugin);
            assert(yolo != nullptr);
            yolo->setName(layerName.c_str());
            std::string inputVol = dimsToString(previous->getDimensions());
            previous = yolo->getOutput(0);
            assert(previous != nullptr);
            previous->setName(layerName.c_str());
            std::string outputVol = dimsToString(previous->getDimensions());
            m_Network->markOutput(*previous);
            channels = getNumChannels(previous);
            tensorOutputs.push_back(yolo->getOutput(0));
            printLayerInfo(layerIndex, "yolo", inputVol, outputVol, std::to_string(weightPtr));
            ++outputTensorCount;
        }
        else if (m_configBlocks.at(i).at("type") == "route")
        {
            size_t found = m_configBlocks.at(i).at("layers").find(",");
            if (found != std::string::npos)//concate multi layers
            {
                std::vector<int> vec_index = split_layer_index(m_configBlocks.at(i).at("layers"), ",");
                for (auto &ind_layer:vec_index)
                {
                    if (ind_layer < 0)
                    {
                        ind_layer = static_cast<int>(tensorOutputs.size()) + ind_layer;
                    }
                    assert(ind_layer < static_cast<int>(tensorOutputs.size()) && ind_layer >= 0);
                }
                nvinfer1::ITensor** concatInputs
                    = reinterpret_cast<nvinfer1::ITensor**>(malloc(sizeof(nvinfer1::ITensor*) * vec_index.size()));
                for (size_t ind = 0; ind < vec_index.size(); ++ind)
                {
                    concatInputs[ind] = tensorOutputs[vec_index[ind]];
                }
                nvinfer1::IConcatenationLayer* concat
                    = m_Network->addConcatenation(concatInputs, static_cast<int>(vec_index.size()));
                assert(concat != nullptr);
                std::string concatLayerName = "route_" + std::to_string(i - 1);
                concat->setName(concatLayerName.c_str());
                // concatenate along the channel dimension
                concat->setAxis(0);
                previous = concat->getOutput(0);
                assert(previous != nullptr);
                nvinfer1::Dims debug = previous->getDimensions();
                std::string outputVol = dimsToString(previous->getDimensions());
                int nums = 0;
                for (auto &indx:vec_index)
                {
                    nums += getNumChannels(tensorOutputs[indx]);
                }
                channels = nums;
                tensorOutputs.push_back(concat->getOutput(0));
                printLayerInfo(layerIndex, "route", "        -", outputVol,std::to_string(weightPtr));
            }
            else //single layer
            {
                int idx = std::stoi(trim(m_configBlocks.at(i).at("layers")));
                if (idx < 0)
                {
                    idx = static_cast<int>(tensorOutputs.size()) + idx;
                }
                assert(idx < static_cast<int>(tensorOutputs.size()) && idx >= 0);

                //route
                if (m_configBlocks.at(i).find("groups") == m_configBlocks.at(i).end())
                {
                    previous = tensorOutputs[idx];
                    assert(previous != nullptr);
                    std::string outputVol = dimsToString(previous->getDimensions());
                    // set the output volume depth
                    channels = getNumChannels(tensorOutputs[idx]);
                    tensorOutputs.push_back(tensorOutputs[idx]);
                    printLayerInfo(layerIndex, "route", "        -", outputVol, std::to_string(weightPtr));

                }
                //yolov4-tiny route split layer
                else
                {
                    if (m_configBlocks.at(i).find("group_id") == m_configBlocks.at(i).end())
                    {
                        assert(0);
                    }
                    int chunk_idx = std::stoi(trim(m_configBlocks.at(i).at("group_id")));
                    nvinfer1::ILayer* out = layer_split(i, tensorOutputs[idx], m_Network);
                    std::string inputVol = dimsToString(previous->getDimensions());
                    previous = out->getOutput(chunk_idx);
                    assert(previous != nullptr);
                    channels = getNumChannels(previous);
                    std::string outputVol = dimsToString(previous->getDimensions());
                    tensorOutputs.push_back(out->getOutput(chunk_idx));
                    printLayerInfo(layerIndex,"chunk", inputVol, outputVol, std::to_string(weightPtr));
                }
            }
        }
        else if (m_configBlocks.at(i).at("type") == "upsample")
        {
            std::string inputVol = dimsToString(previous->getDimensions());
            nvinfer1::ILayer* out = netAddUpsample(i - 1, m_configBlocks[i], weights, trtWeights,
                                                   channels, previous, m_Network);
            previous = out->getOutput(0);
            std::string outputVol = dimsToString(previous->getDimensions());
            tensorOutputs.push_back(out->getOutput(0));
            printLayerInfo(layerIndex, "upsample", inputVol, outputVol, "    -");
        }
        else if (m_configBlocks.at(i).at("type") == "maxpool")
        {
            // Add same padding layers
            if (m_configBlocks.at(i).at("size") == "2" && m_configBlocks.at(i).at("stride") == "1")
            {
                m_TinyMaxpoolPaddingFormula->addSamePaddingLayer("maxpool_" + std::to_string(i));
            }
            std::string inputVol = dimsToString(previous->getDimensions());
            nvinfer1::ILayer* out = netAddMaxpool(i, m_configBlocks.at(i), previous, m_Network);
            previous = out->getOutput(0);
            assert(previous != nullptr);
            std::string outputVol = dimsToString(previous->getDimensions());
            tensorOutputs.push_back(out->getOutput(0));
            printLayerInfo(layerIndex, "maxpool", inputVol, outputVol, std::to_string(weightPtr));
        }
        else
        {
            std::cout << "Unsupported layer type --> \"" << m_configBlocks.at(i).at("type") << "\""
                      << std::endl;
            assert(0);
        }
    }

    if (static_cast<int>(weights.size()) != weightPtr)
    {
        std::cout << "Number of unused weights left : " << static_cast<int>(weights.size()) - weightPtr << std::endl;
        assert(0);
    }

 //   std::cout << "Output blob names :" << std::endl;
 //   for (auto& tensor : m_OutputTensors) std::cout << tensor.blobName << std::endl;

    // Create and cache the engine if not already present
    if (fileExists(m_EnginePath))
    {
        std::cout << "Using previously generated plan file located at " << m_EnginePath
                  << std::endl;
        destroyNetworkUtils(trtWeights);
        return;
    }

    /*std::cout << "Unable to find cached TensorRT engine for network : " << m_NetworkType
              << " precision : " << m_Precision << " and batch size :" << m_BatchSize << std::endl;*/

    m_Builder->setMaxBatchSize(m_BatchSize);
    //m_Builder->setMaxWorkspaceSize(1 << 20);

    config->setMaxWorkspaceSize(1 << 20);
    if (dataType == nvinfer1::DataType::kINT8)
    {
        assert((calibrator != nullptr) && "Invalid calibrator for INT8 precision");
      //  m_Builder->setInt8Mode(true);
        config->setFlag(nvinfer1::BuilderFlag::kINT8);
     //   m_Builder->setInt8Calibrator(calibrator);
        config->setInt8Calibrator(calibrator);
    //  config->setTacticSources(1U << static_cast<uint32_t>(TacticSource::kCUBLAS) | 1U << static_cast<uint32_t>(TacticSource::kCUBLAS_LT));
    }
    else if (dataType == nvinfer1::DataType::kHALF)
    {
        config->setFlag(nvinfer1::BuilderFlag::kFP16);
     //   m_Builder->setHalf2Mode(true);
    }

    m_Builder->allowGPUFallback(true);
    int nbLayers = m_Network->getNbLayers();
    int layersOnDLA = 0;
 //   std::cout << "Total number of layers: " << nbLayers << std::endl;
    for (int i = 0; i < nbLayers; i++)
    {
        nvinfer1::ILayer* curLayer = m_Network->getLayer(i);
        if (m_DeviceType == "kDLA" && m_Builder->canRunOnDLA(curLayer))
        {
            m_Builder->setDeviceType(curLayer, nvinfer1::DeviceType::kDLA);
            layersOnDLA++;
            std::cout << "Set layer " << curLayer->getName() << " to run on DLA" << std::endl;
        }
    }
 //   std::cout << "Total number of layers on DLA: " << layersOnDLA << std::endl;

    // Build the engine
    std::cout << "Building the TensorRT Engine..." << std::endl;
    m_Engine = m_Builder->buildEngineWithConfig(*m_Network,*config);
    assert(m_Engine != nullptr);
    std::cout << "Building complete!" << std::endl;

    // Serialize the engine
    writePlanFileToDisk();

    // destroy
    destroyNetworkUtils(trtWeights);
}

void Detecter::doInference(const unsigned char* input, const uint32_t batchSize)
{
    Timer timer;
    // std::vector<int64_t> bufferSize;
    // int nbBindings = m_Engine->getNbBindings();
    // bufferSize.resize(nbBindings);

    // for (int i = 0; i < nbBindings; ++i)
    // {
    //     nvinfer1::Dims dims = m_Engine->getBindingDimensions(i);
    //     nvinfer1::DataType dtype = m_Engine->getBindingDataType(i);
    //     int64_t totalSize = volume(dims) * 1 * getElementSize(dtype);
    //     bufferSize[i] = totalSize;
    //     std::cout << "binding" << i << ": " << totalSize << std::endl;
    //     cudaMalloc(&s_buffers[i], totalSize);
    // }

    // CHECK(cudaMalloc(&s_buffers[0], batchSize *416 * 416 * 3 * sizeof(float)));
    // CHECK(cudaMalloc(&s_buffers[1], batchSize * 3 * sizeof(float)));
    // CHECK(cudaMalloc(&s_buffers[2], batchSize * 3 * sizeof(float)));
    // CHECK(cudaMalloc(&s_buffers[3], batchSize * 3 * sizeof(float)));

    assert(batchSize <= m_BatchSize && "Image batch size exceeds TRT engines batch size");
    // cudaMemcpyAsync(s_buffers[0], input,
    //                               batchSize * m_InputSize * sizeof(float), cudaMemcpyHostToDevice,
    //                               m_CudaStream);

    // m_Context->enqueue(batchSize, s_buffers, m_CudaStream, nullptr);
    NV_CUDA_CHECK(cudaMemcpyAsync(m_DeviceBuffers.at(m_InputBindingIndex), input,
                                  batchSize * m_InputSize * sizeof(float), cudaMemcpyHostToDevice,
                                  m_CudaStream));

    m_Context->enqueue(batchSize, m_DeviceBuffers.data(), m_CudaStream, nullptr);
    for (auto& tensor : m_OutputTensors)
    {
        NV_CUDA_CHECK(cudaMemcpyAsync(tensor.hostBuffer, m_DeviceBuffers.at(tensor.bindingIndex),
                                      batchSize * tensor.volume * sizeof(float),
                                      cudaMemcpyDeviceToHost, m_CudaStream));
    }


    // int outSize1 = bufferSize[1] / sizeof(float) / BATCH_SIZE;
    // int outSize2 = bufferSize[2] / sizeof(float) / BATCH_SIZE;
    // int outSize3 = bufferSize[3] / sizeof(float) / BATCH_SIZE;
    // std::cout <<  bufferSize[1]  << std::endl;
    // std::cout <<  bufferSize[2]  << std::endl;
    // std::cout <<  bufferSize[3] << std::endl;

    // auto *out1 = new float[batchSize * 3];
    // auto *out2 = new float[batchSize * 3];
    // auto *out3 = new float[batchSize * 3];
    // cudaMemcpyAsync(out1, s_buffers[1], bufferSize[1], cudaMemcpyDeviceToHost, m_CudaStream);
    // cudaMemcpyAsync(out2, s_buffers[2], bufferSize[2], cudaMemcpyDeviceToHost, m_CudaStream);
    // cudaMemcpyAsync(out3, s_buffers[3], bufferSize[3], cudaMemcpyDeviceToHost, m_CudaStream);
    // CHECK(cudaMemcpyAsync(out1, s_buffers[1], batchSize * 3 * sizeof(float), cudaMemcpyDeviceToHost, m_CudaStream));
    // CHECK(cudaMemcpyAsync(out2, s_buffers[2], batchSize * 3 * sizeof(float), cudaMemcpyDeviceToHost, m_CudaStream));
    // CHECK(cudaMemcpyAsync(out3, s_buffers[3], batchSize * 3 * sizeof(float), cudaMemcpyDeviceToHost, m_CudaStream));
    // exit(0);


    cudaStreamSynchronize(m_CudaStream);
    timer.out("inference");
}

std::vector<BBoxInfo> Detecter::decodeDetections(const int& imageIdx,
                                             const int& imageH,
                                             const int& imageW)
{
    Timer timer;
    std::vector<BBoxInfo> binfo;
    for (auto& tensor : m_OutputTensors)
    {
        std::vector<BBoxInfo> curBInfo = decodeTensor(imageIdx, imageH, imageW, tensor);
        binfo.insert(binfo.end(), curBInfo.begin(), curBInfo.end());
    }
    timer.out("decodeDetections");
    return binfo;
}


void Detecter::allocateBuffers()
{
    m_DeviceBuffers.resize(m_Engine->getNbBindings(), nullptr);
    assert(m_InputBindingIndex != -1 && "Invalid input binding index");
    NV_CUDA_CHECK(cudaMalloc(&m_DeviceBuffers.at(m_InputBindingIndex),
                             m_BatchSize * m_InputSize * sizeof(float)));

    for (auto& tensor : m_OutputTensors)
    {
        tensor.bindingIndex = m_Engine->getBindingIndex(tensor.blobName.c_str());
        assert((tensor.bindingIndex != -1) && "Invalid output binding index");
        NV_CUDA_CHECK(cudaMalloc(&m_DeviceBuffers.at(tensor.bindingIndex),
                                 m_BatchSize * tensor.volume * sizeof(float)));
        NV_CUDA_CHECK(
            cudaMallocHost(&tensor.hostBuffer, tensor.volume * m_BatchSize * sizeof(float)));
    }
}

bool Detecter::verifyEngine()
{
    assert((m_Engine->getNbBindings() == (1 + m_OutputTensors.size())
            && "Binding info doesn't match between cfg and engine file \n"));

    for (auto tensor : m_OutputTensors)
    {
        assert(!strcmp(m_Engine->getBindingName(tensor.bindingIndex), tensor.blobName.c_str())
               && "Blobs names dont match between cfg and engine file \n");
        assert(get3DTensorVolume(m_Engine->getBindingDimensions(tensor.bindingIndex))
                   == tensor.volume
               && "Tensor volumes dont match between cfg and engine file \n");
    }

    assert(m_Engine->bindingIsInput(m_InputBindingIndex) && "Incorrect input binding index \n");
    assert(m_Engine->getBindingName(m_InputBindingIndex) == m_InputBlobName
           && "Input blob name doesn't match between config and engine file");
    assert(get3DTensorVolume(m_Engine->getBindingDimensions(m_InputBindingIndex)) == m_InputSize);
    return true;
}

void Detecter::destroyNetworkUtils(std::vector<nvinfer1::Weights>& trtWeights)
{
    if (m_Network) m_Network->destroy();
    if (m_Engine) m_Engine->destroy();
    if (m_Builder) m_Builder->destroy();
    if (m_ModelStream) m_ModelStream->destroy();

    // deallocate the weights
    for (auto & trtWeight : trtWeights)
    {
        if (trtWeight.count > 0) free(const_cast<void*>(trtWeight.values));
    }
}

std::vector<BBoxInfo>  Detecter::decodeTensor(const int imageIdx, const int imageH, const int imageW, const TensorInfo& tensor)
{
    float   scale_h = 1.f;
    float   scale_w = 1.f;
    int xOffset = 0;
    int yOffset = 0;

    const float* detections = &tensor.hostBuffer[imageIdx * tensor.volume];

    std::vector<BBoxInfo> binfo;
    for (uint32_t y = 0; y < tensor.grid_h; ++y)
    {
        for (uint32_t x = 0; x < tensor.grid_w; ++x)
        {
            for (uint32_t b = 0; b < tensor.numBBoxes; ++b)
            {
                const float pw = tensor.anchors[tensor.masks[b] * 2];
                const float ph = tensor.anchors[tensor.masks[b] * 2 + 1];

                const int numGridCells = tensor.grid_h * tensor.grid_w;
                const int bbindex = y * tensor.grid_w + x;
                const float bx
                    = x + detections[bbindex + numGridCells * (b * (5 + tensor.numClasses) + 0)];
                const float by
                    = y + detections[bbindex + numGridCells * (b * (5 + tensor.numClasses) + 1)];
                const float bw
                    = pw * detections[bbindex + numGridCells * (b * (5 + tensor.numClasses) + 2)];
                const float bh
                    = ph * detections[bbindex + numGridCells * (b * (5 + tensor.numClasses) + 3)];

                const float objectness
                    = detections[bbindex + numGridCells * (b * (5 + tensor.numClasses) + 4)];

                float maxProb = 0.0f;
                int maxIndex = -1;

                for (uint32_t i = 0; i < tensor.numClasses; ++i)
                {
                    float prob
                        = (detections[bbindex
                            + numGridCells * (b * (5 + tensor.numClasses) + (5 + i))]);

                    if (prob > maxProb)
                    {
                        maxProb = prob;
                        maxIndex = i;
                    }
                }
                maxProb = objectness * maxProb;
                // printf("-=-=-=-=-=-%f\n",maxProb);

                if (maxProb > m_ProbThresh)
                {
                    add_bbox_proposal(bx, by, bw, bh, tensor.stride_h, tensor.stride_w, scale_h, scale_w, xOffset, yOffset, maxIndex, maxProb, imageW, imageH, binfo);
                }
            }
        }
    }
    return binfo;
}


void Detecter::setOutput(int type)
{
    m_OutputTensors.clear();
    printf("0-0-0-0-0-0------------------%d",type);
    if(type==2)
        for (int i = 0; i < 2; ++i)
        {
            TensorInfo outputTensor;
            outputTensor.numClasses = CLASS_BUM;
            outputTensor.blobName = "yolo_" + std::to_string(i);
            outputTensor.gridSize = (m_InputH / 32) * pow(2, i);
            outputTensor.grid_h = (m_InputH / 32) * pow(2, i);
            outputTensor.grid_w = (m_InputW / 32) * pow(2, i);
            outputTensor.stride = m_InputH / outputTensor.gridSize;
            outputTensor.stride_h = m_InputH / outputTensor.grid_h;
            outputTensor.stride_w = m_InputW / outputTensor.grid_w;
            outputTensor.numBBoxes = 3;
            outputTensor.volume = outputTensor.grid_h* outputTensor.grid_w
                *(outputTensor.numBBoxes*(5 + outputTensor.numClasses));
            if(i==1)
            {
                outputTensor.masks.push_back(1);
                outputTensor.masks.push_back(2);
                outputTensor.masks.push_back(3);
            }
            if(i==0)
            {
                outputTensor.masks.push_back(3);
                outputTensor.masks.push_back(4);
                outputTensor.masks.push_back(5);
            }
            outputTensor.anchors.push_back(10);
            outputTensor.anchors.push_back(14);
            outputTensor.anchors.push_back(23);
            outputTensor.anchors.push_back(27);
            outputTensor.anchors.push_back(37);
            outputTensor.anchors.push_back(58);
            outputTensor.anchors.push_back(81);
            outputTensor.anchors.push_back(82);
            outputTensor.anchors.push_back(135);
            outputTensor.anchors.push_back(169);
            outputTensor.anchors.push_back(344);
            outputTensor.anchors.push_back(319);

            if (m_ClassNames.empty())
            {
                for (uint32_t j=0;j< outputTensor.numClasses;++j)
                {
                    m_ClassNames.push_back(std::to_string(j));
                }
            }
            m_OutputTensors.push_back(outputTensor);
        }
    else
        for (int i = 0; i < 3; ++i)
        {
            TensorInfo outputTensor;
            outputTensor.numClasses = CLASS_BUM;
            outputTensor.blobName = "yolo_" + to_string(i);
            // if (i==0)
            // {
            //     outputTensor.blobName = "139_convolutional_reshape_2";
            // }else if (i==1)
            // {
            //     outputTensor.blobName = "150_convolutional_reshape_2";
            // }else if (i==2)
            // {
            //     outputTensor.blobName = "161_convolutional_reshape_2";
            // }
            outputTensor.gridSize = (m_InputH / 32) * pow(2, 2-i);
            outputTensor.grid_h = (m_InputH / 32) * pow(2, 2-i);
            outputTensor.grid_w = (m_InputW / 32) * pow(2, 2-i);
            outputTensor.stride = m_InputH / outputTensor.gridSize;
            outputTensor.stride_h = m_InputH / outputTensor.grid_h;
            outputTensor.stride_w = m_InputW / outputTensor.grid_w;
            outputTensor.numBBoxes = 3;
            outputTensor.volume = outputTensor.grid_h* outputTensor.grid_w
                *(outputTensor.numBBoxes*(5 + outputTensor.numClasses));
            if(i==0)
            {
                outputTensor.masks.push_back(0);
                outputTensor.masks.push_back(1);
                outputTensor.masks.push_back(2);
            }
            if(i==1)
            {
                outputTensor.masks.push_back(3);
                outputTensor.masks.push_back(4);
                outputTensor.masks.push_back(5);
            }
            if(i==2)
            {
                outputTensor.masks.push_back(6);
                outputTensor.masks.push_back(7);
                outputTensor.masks.push_back(8);
            }
            outputTensor.anchors.push_back(12);
            outputTensor.anchors.push_back(16);
            outputTensor.anchors.push_back(19);
            outputTensor.anchors.push_back(36);
            outputTensor.anchors.push_back(40);
            outputTensor.anchors.push_back(28);
            outputTensor.anchors.push_back(36);
            outputTensor.anchors.push_back(75);
            outputTensor.anchors.push_back(76);
            outputTensor.anchors.push_back(55);
            outputTensor.anchors.push_back(72);
            outputTensor.anchors.push_back(146);
            outputTensor.anchors.push_back(142);
            outputTensor.anchors.push_back(110);
            outputTensor.anchors.push_back(192);
            outputTensor.anchors.push_back(243);
            outputTensor.anchors.push_back(459);
            outputTensor.anchors.push_back(401);

            if (m_ClassNames.empty())
            {
                for (uint32_t j=0;j< outputTensor.numClasses;++j)
                {
                    m_ClassNames.push_back(std::to_string(j));
                }
            }
            m_OutputTensors.push_back(outputTensor);
        }
}

void Detecter::writePlanFileToDisk()
{
    std::cout << "Serializing the TensorRT Engine..." << std::endl;
    assert(m_Engine && "Invalid TensorRT Engine");
    m_ModelStream = m_Engine->serialize();
    assert(m_ModelStream && "Unable to serialize engine");
    assert(!m_EnginePath.empty() && "Enginepath is empty");

    // write data to output file
    std::stringstream gieModelStream;
    gieModelStream.seekg(0, gieModelStream.beg);
    gieModelStream.write(static_cast<const char*>(m_ModelStream->data()), m_ModelStream->size());
    std::ofstream outFile;
    outFile.open(m_EnginePath, std::ios::binary | std::ios::out);
    outFile << gieModelStream.rdbuf();
    outFile.close();

    std::cout << "Serialized plan file cached at location : " << m_EnginePath << std::endl;
}