#pragma once

#include <fstream>
#include <vector>
#include <string>
#include <random>
#include <filesystem>
#include <opencv2/imgproc.hpp>
#include <opencv2/opencv.hpp>
#include <opencv2/dnn.hpp>

class Yolo
{
public:
	struct Detection
	{
		int class_id = 0;
		std::string className;
		float confidence = 0.0;
		int priority = -1;
		cv::Scalar color;
		cv::Rect box;
	};

private:
	static constexpr float modelConfidenceThreshold = 0.25;
	static constexpr float modelScoreThreshold = 0.45;
	static constexpr float modelNMSThreshold = 0.50;

	void loadClasses(const std::string& classes);
	void loadOnnxNetwork(const std::filesystem::path& path);
	cv::Mat formatToSquare(const cv::Mat &source);

	std::string modelPath;

	std::vector<std::pair<std::string, int>> classes;

	cv::Size2f modelShape;

	bool letterBoxForSquare = true;

	cv::dnn::Net net;

public:
	Yolo(const std::filesystem::path &onnxModelPath = "", const cv::Size& modelInputShape = {640, 480},
		const std::filesystem::path& classesTxtFilePath = "", bool runWithOCl = true);
	std::vector<Detection> runInference(const cv::Mat &input);
	int getClassForStr(const std::string& str) const;
};