#include "trt_utils.h"
#include <NvInferRuntimeCommon.h>
#include <experimental/filesystem>
#include <fstream>
#include <iomanip>
using namespace nvinfer1;
REGISTER_TENSORRT_PLUGIN(MishPluginCreator);
REGISTER_TENSORRT_PLUGIN(ChunkPluginCreator);
REGISTER_TENSORRT_PLUGIN(HardswishPluginCreator);

cv::Mat blobFromDsImages(const std::vector<DsImage>& inputImages,
                        const int& inputH,
                         const int& inputW)
{
    std::vector<cv::Mat> letterboxStack(inputImages.size());
    for (uint32_t i = 0; i < inputImages.size(); ++i)
    {
        inputImages.at(i).getLetterBoxedImage().copyTo(letterboxStack.at(i));
    }
    return cv::dnn::blobFromImages(letterboxStack, 1.0, cv::Size(inputW, inputH),
                                   cv::Scalar(0.0, 0.0, 0.0),true);
}

// static void leftTrim(std::string& s)
// {
//     s.erase(s.begin(), find_if(s.begin(), s.end(), [](int ch) { return !isspace(ch); }));
// }

// static void rightTrim(std::string& s)
// {
//     s.erase(find_if(s.rbegin(), s.rend(), [](int ch) { return !isspace(ch); }).base(), s.end());
// }

// std::string trim(std::string s)
// {
//     leftTrim(s);
//     rightTrim(s);
//     return s;
// }

// std::string triml(std::string s,const char* t)
// {
//     s.erase(0, s.find_first_not_of(t));
//     return s;
// }

// std::string trimr(std::string s, const char* t)
// {
//     s.erase(s.find_last_not_of(t) + 1);
//     return s;
// }

float clamp(const float val, const float minVal, const float maxVal)
{
    assert(minVal <= maxVal);
    return std::min(maxVal, std::max(minVal, val));
}

bool fileExists(const std::string fileName, bool verbose)
{
    if (!std::experimental::filesystem::exists(std::experimental::filesystem::path(fileName)))
    {
        if (verbose) std::cout << "File does not exist : " << fileName << std::endl;
        return false;
    }
    return true;
}

// BBox convertBBoxNetRes(const float& bx, const float& by, const float& bw, const float& bh,
//                        const uint32_t& stride, const uint32_t& netW, const uint32_t& netH)
// {
//     BBox b;
//     // Restore coordinates to network input resolution
//     float x = bx * stride;
//     float y = by * stride;

//     b.x1 = x - bw / 2;
//     b.x2 = x + bw / 2;

//     b.y1 = y - bh / 2;
//     b.y2 = y + bh / 2;

//     b.x1 = clamp(b.x1, 0, netW);
//     b.x2 = clamp(b.x2, 0, netW);
//     b.y1 = clamp(b.y1, 0, netH);
//     b.y2 = clamp(b.y2, 0, netH);

//     return b;
// }

// void convertBBoxImgRes(const float scalingFactor,
//     const float xOffset,
//     const float yOffset,
//     BBox& bbox)
// {
//         //// Undo Letterbox
//     bbox.x1 -= xOffset;
//     bbox.x2 -= xOffset;
//     bbox.y1 -= yOffset;
//     bbox.y2 -= yOffset;
// //// Restore to input resolution
//     bbox.x1 /= scalingFactor;
//     bbox.x2 /= scalingFactor;
//     bbox.y1 /= scalingFactor;
//     bbox.y2 /= scalingFactor;
//     std::cout << "convertBBoxImgRes" << std::endl;


// }

// void printPredictions(const BBoxInfo& b, const std::string& className)
// {
//     std::cout << " label:" << b.label << "(" << className << ")"
//               << " confidence:" << b.prob << " xmin:" << b.box.x1 << " ymin:" << b.box.y1
//               << " xmax:" << b.box.x2 << " ymax:" << b.box.y2 << std::endl;
// }
//
uint64_t get3DTensorVolume(nvinfer1::Dims inputDims)
{
    assert(inputDims.nbDims == 3);
    return inputDims.d[0] * inputDims.d[1] * inputDims.d[2];
}


nvinfer1::ICudaEngine* loadTRTEngine(const std::string planFilePath, PluginFactory* pluginFactory,
                                     Logger& logger)
{
    // reading the model in memory
    std::cout << "Loading TRT Engine..." << std::endl;
    assert(fileExists(planFilePath));
    std::stringstream trtModelStream;
    trtModelStream.seekg(0, trtModelStream.beg);
    std::ifstream cache(planFilePath,std::ios::binary | std::ios::in);
    assert(cache.good());
    trtModelStream << cache.rdbuf();
    cache.close();

    // calculating model size
    trtModelStream.seekg(0, std::ios::end);
    const auto modelSize = trtModelStream.tellg();
    trtModelStream.seekg(0, std::ios::beg);
    void* modelMem = malloc(modelSize);
    trtModelStream.read((char*) modelMem, modelSize);

    nvinfer1::IRuntime* runtime = nvinfer1::createInferRuntime(logger);
    std::cout << "test................................" << std::endl;
    nvinfer1::ICudaEngine* engine
        = runtime->deserializeCudaEngine(modelMem, modelSize, pluginFactory);
    free(modelMem);
    runtime->destroy();
    std::cout << "Loading Complete!" << std::endl;

    return engine;
}

std::vector<BBoxInfo> nmsAllClasses(const float nmsThresh,
    std::vector<BBoxInfo>& binfo,
    const uint32_t numClasses,
    const std::string &model_type)
{
    std::vector<BBoxInfo> result;
    std::vector<std::vector<BBoxInfo>> splitBoxes(numClasses);
    for (auto& box : binfo)
    {
        splitBoxes.at(box.label).push_back(box);
    }

    for (auto& boxes : splitBoxes)
    {
        boxes = nonMaximumSuppression(nmsThresh, boxes);
        result.insert(result.end(), boxes.begin(), boxes.end());
    }

    return result;
}

std::vector<BBoxInfo> nonMaximumSuppression(const float nmsThresh, std::vector<BBoxInfo> binfo)
{
    auto overlap1D = [](float x1min, float x1max, float x2min, float x2max) -> float
    {
        if (x1min > x2min)
        {
            std::swap(x1min, x2min);
            std::swap(x1max, x2max);
        }
        return x1max < x2min ? 0 : std::min(x1max, x2max) - x2min;
    };
    auto computeIoU = [&overlap1D](BBox& bbox1, BBox& bbox2) -> float
    {
        float overlapX = overlap1D(bbox1.x1, bbox1.x2, bbox2.x1, bbox2.x2);
        float overlapY = overlap1D(bbox1.y1, bbox1.y2, bbox2.y1, bbox2.y2);
        float area1 = (bbox1.x2 - bbox1.x1) * (bbox1.y2 - bbox1.y1);
        float area2 = (bbox2.x2 - bbox2.x1) * (bbox2.y2 - bbox2.y1);
        float overlap2D = overlapX * overlapY;
        float u = area1 + area2 - overlap2D;
        return u == 0 ? 0 : overlap2D / u;
    };

    std::stable_sort(binfo.begin(), binfo.end(),
                     [](const BBoxInfo& b1, const BBoxInfo& b2) { return b1.prob > b2.prob; });
    std::vector<BBoxInfo> out;
    for (auto& i : binfo)
    {
        bool keep = true;
        for (auto& j : out)
        {
            if (keep)
            {
                float overlap = computeIoU(i.box, j.box);
                keep = overlap <= nmsThresh;
            }
            else
                break;
        }
        if (keep) out.push_back(i);
    }
    return out;
}