#include <regex>
#include <iostream>
#include <filesystem>
#include <ctime>
#include "utils.h"
#include "yolov8Predictor.h"
#include "opencv2/opencv.hpp"

int main(int argc, char *argv[])
{
    // yolov8参数
    float confThreshold = 0.4f;
    float iouThreshold = 0.4f;
    float maskThreshold = 0.5f;
    bool isGPU = true;

    // 模型路径
    const std::string classNamePath = "../model/coco.names";
    const std::vector<std::string> classNames = utils::loadNames(classNamePath);
    const std::string modelPath = "../model/yolov8n.onnx";
    //输入图像路径 
    const std::string imagePath = "../img_in/";
    const std::string _strPattern = imagePath + "*.jpg"; 
    std::vector<cv::String> ImageNames;
    cv::glob(_strPattern, ImageNames);
    //输出图像路径
    const std::string savePath = "../img_out";
    const std::string suffixName = "png";

    if (classNames.empty())
    {
        std::cerr << "Error: Empty class names file." << std::endl;
        return -1;
    }
    if (!std::filesystem::exists(modelPath))
    {
        std::cerr << "Error: There is no model." << std::endl;
        return -1;
    }
    if (!std::filesystem::is_directory(imagePath))
    {
        std::cerr << "Error: There is no img." << std::endl;
        return -1;
    }
    if (!std::filesystem::is_directory(savePath))
    {
        std::filesystem::create_directory(savePath);
    }
    std::cout << "Model from :::" << modelPath << std::endl;
    std::cout << "Images from :::" << imagePath << std::endl;
    std::cout << "Resluts will be saved :::" << savePath << std::endl;

    YOLOPredictor predictor{nullptr};
    try
    {
        predictor = YOLOPredictor(modelPath, isGPU,
                                  confThreshold,
                                  iouThreshold,
                                  maskThreshold);
        std::cout << "Model was initialized." << std::endl;
    }
    catch (const std::exception &e)
    {
        std::cerr << e.what() << std::endl;
        return -1;
    }
    assert(classNames.size() == predictor.classNums);
    std::regex pattern(".+\\.(jpg|jpeg|png|gif)$");
    std::cout << "Start predicting..." << std::endl;

    clock_t startTime, endTime;
    startTime = clock();

    int picNums = 0;
    int batch_size = 4;
    std::vector<cv::Mat> imagebatch;
    for(int i = 0; i < batch_size; i++)
    {
        std::cout << "image name:" << ImageNames[i] << std::endl;
        cv::Mat image = cv::imread(ImageNames[i]);
        if (image.empty())
        {
            std::cerr << "Error: Empty image." << std::endl;
            return -1;
        }
        imagebatch.push_back(image);

    }
    std::vector<std::vector<Yolov8Result>> results = predictor.predict(imagebatch);
    for(int i = 0; i < results.size(); i++)
    {
        std::vector<Yolov8Result> result = results[i];
        cv::Mat image = imagebatch[i];
        utils::visualizeDetection(image, result, classNames);
        std::string newFilename = std::filesystem::path(ImageNames[i]).filename().string();
        std::string outputFilename = savePath + "/" + newFilename;
        cv::imwrite(outputFilename, image);
        std::cout << "Saved !!!" << std::endl;
    }
    endTime = clock();
    std::cout << "The total run time is: " << (double)(endTime - startTime) / CLOCKS_PER_SEC << "seconds" << std::endl;
    std::cout << "The average run time is: " << (double)(endTime - startTime) / batch_size / CLOCKS_PER_SEC << "seconds" << std::endl;
    std::cout << "##########DONE################" << std::endl;
    
    return 0;
}