#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <SimpleLog.h>
#include <Filesystem.h>
#include <Classifier.h>
#include <iostream>
#include <fstream>

int main()
{
    // 创建分类器
    ortSamples::Classifier classifier;
    ortSamples::InitializationParameterOfClassifier initParamOfClassifier;
    initParamOfClassifier.configFilePath=CONFIG_FILE;
    ortSamples::ErrorCode errorCode=classifier.Initialize(initParamOfClassifier);
    if(errorCode!=ortSamples::SUCCESS)
    {
        LOG_ERROR(stdout, "fail to initialize ResNet50!\n");
        exit(-1);
    }
    LOG_INFO(stdout, "succeed to initialize ResNet50\n");

    // 读取测试图片
    cv::Mat srcImage=cv::imread("../Resource/Images/ImageNet_01.jpg",1);

    // 设置batchsize
    int batchsize=1;
    std::vector<cv::Mat> srcImages;
    for(int i=0;i<batchsize;++i)
    {
        srcImages.push_back(srcImage);
    }

    // 推理
    std::vector<std::vector<ortSamples::ResultOfPrediction>> predictions;
    double time1 = cv::getTickCount();
    classifier.Classify(srcImages,predictions);
    double time2 = cv::getTickCount();
    double elapsedTime = (time2 - time1)*1000 / cv::getTickFrequency();
    LOG_INFO(stdout, "inference time:%f ms\n", elapsedTime);
    // 读取分类列表
    std::ifstream file("../Resource/synset.txt");
    if (!file.is_open()) {
        std::cerr << "无法打开synset.txt文件" << std::endl;
        return 1;
    }

    
    // 获取推理结果
    LOG_INFO(stdout,"========== Classification Results ==========\n");
    for(int i=0;i<predictions.size();++i)
    {
        // 一个batch中第i幅图像的结果
        LOG_INFO(stdout,"========== %d result ==========\n",i);
        std::vector<ortSamples::ResultOfPrediction> resultOfPredictions=predictions[i];
        double maxConfidence = -DBL_MAX;
        int maxLabel = -1; 
        for(int j=0;j<resultOfPredictions.size();++j)
        {
            ortSamples::ResultOfPrediction prediction=resultOfPredictions[j];
            if (prediction.confidence > maxConfidence) {
                maxConfidence = prediction.confidence;
                maxLabel = prediction.label;
            }
        }
        std::string line;
        int currentLine = 0;
        while (currentLine < maxLabel+1 && std::getline(file, line)) {
            currentLine++;
        }
        LOG_INFO(stdout, "class=%s ; probability=%f\n", line.c_str(), maxConfidence);
        
        //保存图像
        std::string text = std::string("class=") + line.c_str();
        cv::Point text_position(5,20);
        int font_face = cv::FONT_HERSHEY_SIMPLEX;
        double font_scale = 0.5; // 调整字体大小
        cv::Scalar font_color(0, 0, 255); // 白色
        int font_thickness = 1;
        cv::putText(srcImage, text, text_position, font_face, font_scale, font_color, font_thickness);
        cv::imwrite("./output_image.jpg", srcImage);

    }
    file.close();
    return 0;
}
