#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_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)
{
    setOutput(type);
    m_EnginePath = m_staticStruct::model_path;
    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();
}

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();
    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_" + std::to_string(i);
            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);
        }
}