// https://github.com/AlexeyAB/darknet/wiki/How-to-evaluate-accuracy-and-speed-of-YOLOv4
// g++ -I/usr/local/include/opencv4/ main.cpp -lopencv_core -lopencv_imgproc -lopencv_dnn -lopencv_imgcodecs -O3 -std=c++17 -lstdc++fs

#include <iostream>
#include <queue>
#include <iterator>
#include <sstream>
#include <fstream>
#include <iomanip>
#include <chrono>
#include <experimental/filesystem>
namespace fs = std::experimental::filesystem;

#include <opencv2/core.hpp>
#include <opencv2/dnn.hpp>
#include <opencv2/dnn/all_layers.hpp>

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

constexpr float CONFIDENCE_THRESHOLD = 0.001;
constexpr float NMS_THRESHOLD = 0.45;
constexpr int NUM_CLASSES = 80;

int main()
{
    std::vector<std::string> images;
    {
        std::ifstream list("list");
        if (!list)
        {
            std::cerr << "failed to open list\n";
            return 0;
        }

        std::string line;
        while (std::getline(list, line))
            images.push_back(line);
    }

    auto net = cv::dnn::readNetFromDarknet("yolov4.cfg", "yolov4.weights");
    for (auto name : net.getUnconnectedOutLayersNames())
    {
        int layerId = net.getLayerId(name);
        auto layer = net.getLayer(layerId).dynamicCast<cv::dnn::RegionLayer>();
        if (!layer.empty())
            layer->nmsThreshold = 0;
    }

    net.setPreferableBackend(cv::dnn::DNN_BACKEND_CUDA);
    net.setPreferableTarget(cv::dnn::DNN_TARGET_CUDA_FP16);
    auto output_names = net.getUnconnectedOutLayersNames();

    cv::Mat blob;
    std::vector<cv::Mat> detections;

    std::ofstream output("result.json");
    if (!output)
    {
        std::cerr << "failed to open result.json\n";
        return 0;
    }

    output << "[\n";

    for (int idx = 0; idx < images.size(); idx++)
    {
        auto frame = cv::imread(images[idx]);
        if (frame.empty())
        {
            std::cerr << "failed to open " << images[idx] << std::endl;
            continue;
        }

        auto image_id = std::atoll(fs::path(images[idx]).stem().c_str());

        auto total_start = std::chrono::steady_clock::now();
        cv::dnn::blobFromImage(frame, blob, 0.00392, cv::Size(608, 608), cv::Scalar(), true, false, CV_32F);
        net.setInput(blob);

        net.forward(detections, output_names);

        std::vector<int> indices[NUM_CLASSES];
        std::vector<cv::Rect2f> boxes[NUM_CLASSES];
        std::vector<float> scores[NUM_CLASSES];

        for (auto& output : detections)
        {
            const auto num_boxes = output.rows;
            for (int i = 0; i < num_boxes; i++)
            {
                auto x = output.at<float>(i, 0) * frame.cols;
                auto y = output.at<float>(i, 1) * frame.rows;
                auto width = output.at<float>(i, 2) * frame.cols;
                auto height = output.at<float>(i, 3) * frame.rows;

                auto xmin = x - width/2;
                auto ymin = y - height/2;
                auto xmax = x + width/2;
                auto ymax = y + height/2;

                if (xmin < 0) xmin = 0;
                if (ymin < 0) ymin = 0;
                if (xmax > frame.cols) xmax = frame.cols;
                if (ymax > frame.rows) ymax = frame.rows;

                cv::Rect2f rect(xmin, ymin, xmax - xmin, ymax - ymin);

                for (int c = 0; c < NUM_CLASSES; c++)
                {
                    auto confidence = *output.ptr<float>(i, 5 + c);
                    if (confidence >= CONFIDENCE_THRESHOLD)
                    {
                        boxes[c].push_back(rect);
                        scores[c].push_back(confidence);
                    }
                }
            }
        }

        for (int c = 0; c < NUM_CLASSES; c++)
            cv::dnn::NMSBoxes(std::vector<cv::Rect>(std::begin(boxes[c]), std::end(boxes[c])), scores[c], 0.0, NMS_THRESHOLD, indices[c]);

        for (int c= 0; c < NUM_CLASSES; c++)
        {
            for (size_t i = 0; i < indices[c].size(); ++i)
            {
                auto k = indices[c][i];
                const auto rect = boxes[c][k];
                const auto score = scores[c][k];

                if (score)
                {
                    const int coco_ids[] = {1,2,3,4,5,6,7,8,9,10,11,13,14,15,16,17,18,19,20,21,22,23,24,25,27,28,31,32,33,34,35,36,37,38,39,40,41,42,43,44,46,47,48,49,50,51,52,53,54,55,56,57,58,59,60,61,62,63,64,65,67,70,72,73,74,75,76,77,78,79,80,81,82,84,85,86,87,88,89,90};
                    float x = rect.x; 
                    float w = rect.width;
                    float y = rect.y;
                    float h = rect.height;
                    output << "{\"image_id\":" << image_id << ", \"category_id\":" << coco_ids[c] << ", \"bbox\":[" << x << ", " << y << ", " << w << ", " << h << "], \"score\":" << score << "},\n";   
                }
            }
        }
    }

    output.seekp(-2, std::ios_base::cur);
    output << "\n]";

    return 0;
}