#include <thread>
#include <fstream>

#include <opencv2/imgproc.hpp>
#include <opencv2/dnn.hpp>
#include <opencv2/core.hpp>
#include <opencv2/highgui.hpp>

int main(int argc, char** argv)
{
    auto imageToTest = argv[1];
    auto modelToTest = argv[2];
    int modelInputWidth = atoi(argv[3]);
    int modelInputHeight = atoi(argv[4]);

    cv::Size currSize = cv::Size(modelInputWidth, modelInputHeight);
    std::string modelToTestOnnx = modelToTest;
    std::string imagefilename = imageToTest;
    unsigned int num_inferences = 100;

    cv::dnn::Net net = cv::dnn::readNetFromONNX(modelToTestOnnx);

    net.setPreferableBackend(cv::dnn::DNN_BACKEND_CUDA);
    net.setPreferableTarget(cv::dnn::DNN_TARGET_CUDA);

    cv::Mat img = cv::imread(imagefilename, cv::IMREAD_ANYCOLOR);
    cv::Mat resized;
    cv::resize(img, resized, currSize);
	
	std::vector<cv::Mat> imgBatch = { resized };
    bool swaprbchannels = false;
    cv::Mat blob = cv::dnn::blobFromImages(imgBatch, 1.0f / 255.0f, cv::Size(), cv::Scalar(), swaprbchannels, false, CV_32F);

    net.setInput(blob);

    std::vector<cv::String> unconnectedOutLayerNames = net.getUnconnectedOutLayersNames();

    std::vector<cv::Mat> outputs;
    outputs.clear();

    auto timeLoadModelPlusInference1 = std::chrono::high_resolution_clock::now();

    net.forward(outputs, unconnectedOutLayerNames);

    auto timeLoadModelPlusInference2 = std::chrono::high_resolution_clock::now();

    std::chrono::duration<double, std::milli> ms_doubleTimeLoadModelPlusInference = timeLoadModelPlusInference2 - timeLoadModelPlusInference1;

    std::cout << "Execution time (load model + inference): " << ms_doubleTimeLoadModelPlusInference.count() << std::endl; // in ms

    auto time1 = std::chrono::high_resolution_clock::now();

    try {
        for (size_t i = 0; i < num_inferences; i++)
		net.forward(outputs, unconnectedOutLayerNames);
    }
    catch (std::exception& ex)
    {
        std::cout << ex.what() << std::endl;
    }

    auto time2 = std::chrono::high_resolution_clock::now();

    std::chrono::duration<double, std::milli> ms_double = time2 - time1;
    std::cout << "Execution time inference only: " << ms_double.count() / num_inferences << std::endl; // in ms

    std::cout << "Outputs Size: " << outputs[0].size[0] << "x" << outputs[0].size[1] << std::endl;
    std::cout << "Outputs value: " << outputs[0] << std::endl;
}