#include <Filesystem.h>
#include <SimpleLog.h>
#include <YOLOX.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>

void MIGraphXSamplesUsage(char *programName) {
    printf("Usage : %s <index> \n", programName);
    printf("index:\n");
    printf("\t 0) YOLOX sample.\n");
}

void Sample_YOLOX();

int main(int argc, char *argv[]) {
    if (argc < 2 || argc > 2) {
        MIGraphXSamplesUsage(argv[0]);
        return -1;
    }
    if (!strncmp(argv[1], "-h", 2)) {
        MIGraphXSamplesUsage(argv[0]);
        return 0;
    }
    switch (*argv[1]) {
    case '0': {
        Sample_YOLOX();
        break;
    }
    case '1': {
        break;
    }
    default: {
        MIGraphXSamplesUsage(argv[0]);
        break;
    }
    }
    return 0;
}

void Sample_YOLOX() {
    // 创建YOLOX检测器
    migraphxSamples::DetectorYOLOX detector;
    migraphxSamples::InitializationParameterOfDetector initParamOfDetectorYOLOX;
    initParamOfDetectorYOLOX.configFilePath = CONFIG_FILE;
    migraphxSamples::ErrorCode errorCode =
        detector.Initialize(initParamOfDetectorYOLOX, false);
    if (errorCode != migraphxSamples::SUCCESS) {
        LOG_ERROR(stdout, "fail to initialize detector!\n");
        exit(-1);
    }
    LOG_INFO(stdout, "succeed to initialize detector\n");

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

    // 静态推理固定尺寸
    std::vector<std::size_t> inputShape = {1, 3, 640, 640};

    // 推理
    std::vector<migraphxSamples::ResultOfDetection> predictions;
    double time1 = cv::getTickCount();
    detector.Detect(srcImage, inputShape, predictions, false);
    double time2 = cv::getTickCount();
    double elapsedTime = (time2 - time1) * 1000 / cv::getTickFrequency();
    LOG_INFO(stdout, "inference time:%f ms\n", elapsedTime);

    // 获取推理结果
    LOG_INFO(stdout, "========== Detection Results ==========\n");
    for (int i = 0; i < predictions.size(); ++i) {
        migraphxSamples::ResultOfDetection result = predictions[i];
        cv::rectangle(srcImage, result.boundingBox, cv::Scalar(0, 255, 255), 2);

        std::string label = cv::format("%.2f", result.confidence);
        label = result.className + " " + label;
        int left = predictions[i].boundingBox.x;
        int top = predictions[i].boundingBox.y;
        int baseLine;
        cv::Size labelSize =
            cv::getTextSize(label, cv::FONT_HERSHEY_SIMPLEX, 0.5, 1, &baseLine);
        top = max(top, labelSize.height);
        cv::putText(srcImage, label, cv::Point(left, top - 10),
                    cv::FONT_HERSHEY_SIMPLEX, 1, cv::Scalar(0, 255, 255), 2);

        LOG_INFO(stdout, "box:%d %d %d %d,label:%d,confidence:%f\n",
                 predictions[i].boundingBox.x, predictions[i].boundingBox.y,
                 predictions[i].boundingBox.width,
                 predictions[i].boundingBox.height, predictions[i].classID,
                 predictions[i].confidence);
    }
    cv::imwrite("Result.jpg", srcImage);
    LOG_INFO(stdout, "Detection results have been saved to ./Result.jpg\n");
}