#include "YOLOv5Detector.h"

static const vector<string> class_name = { "person", "bicycle", "car", "motorcycle", "airplane", "bus", "train", "truck", "boat", "traffic light",
                      "fire hydrant", "stop sign", "parking meter", "bench", "bird", "cat", "dog", "horse", "sheep", "cow",
                      "elephant", "bear", "zebra", "giraffe", "backpack", "umbrella", "handbag", "tie", "suitcase", "frisbee",
                      "skis", "snowboard", "sports ball", "kite", "baseball bat", "baseball glove", "skateboard", "surfboard",
                      "tennis racket", "bottle", "wine glass", "cup", "fork", "knife", "spoon", "bowl", "banana", "apple",
                      "sandwich", "orange", "broccoli", "carrot", "hot dog", "pizza", "donut", "cake", "chair", "couch",
                      "potted plant", "bed", "dining table", "toilet", "tv", "laptop", "mouse", "remote", "keyboard", "cell phone",
                      "microwave", "oven", "toaster", "sink", "refrigerator", "book", "clock", "vase", "scissors", "teddy bear",
                      "hair drier", "toothbrush" };

YOLOv5Detector::YOLOv5Detector(const wchar_t* model_path)
    : env(ORT_LOGGING_LEVEL_WARNING, "ONNXRuntime"),
    memory_info(Ort::MemoryInfo::CreateCpu(OrtDeviceAllocator, OrtMemTypeCPU)),
    session_options(),
    session(env, model_path, session_options)
{
    // 配置会话选项
    session_options.SetIntraOpNumThreads(5);
    session_options.SetGraphOptimizationLevel(ORT_ENABLE_ALL);

    // 尝试添加 CUDA 提供程序
    try {
        OrtCUDAProviderOptions cuda_options;
        cuda_options.device_id = 0; // 设置为 GPU 0
        session_options.AppendExecutionProvider_CUDA(cuda_options);
        use_cuda = true;
    }
    catch (const Ort::Exception& e) {
        cout << "CUDA 不可用, 使用 CPU." << endl;
        use_cuda = false;
    }

    // 输出是否使用 CUDA
    cout << (use_cuda ? "正在使用 CUDA 进行推理" : "正在使用 CPU 进行推理") << endl;
}

vector<vector<float>> YOLOv5Detector::detect(const cv::Mat& img)
{
    cv::Mat resized_img;
    cv::resize(img, resized_img, cv::Size(640, 640));
    cv::Mat blob = cv::dnn::blobFromImage(resized_img, 1.0 / 255.0, cv::Size(640, 640), cv::Scalar(), true);

    // 定义输入Tensor
    std::array<int64_t, 4> input_shape = { 1, 3, 640, 640 };
    Ort::Value input_tensor = Ort::Value::CreateTensor<float>(memory_info, blob.ptr<float>(), blob.total(), input_shape.data(), input_shape.size());

    // 定义输出Tensor的形状
    std::array<int64_t, 3> output_shape{ 1, 25200, 85 };
    std::vector<float> output_tensor_values(1 * 25200 * 85);
    Ort::Value output_tensor = Ort::Value::CreateTensor<float>(memory_info, output_tensor_values.data(), output_tensor_values.size(), output_shape.data(), output_shape.size());

    // 推理
    session.Run(Ort::RunOptions{ nullptr }, input_names, &input_tensor, 1, output_names, &output_tensor, 1);

    // 处理输出
    float* output_data = output_tensor.GetTensorMutableData<float>();
    int output_size = 25200 * 85;

    vector<vector<float>> info = get_info(output_data, output_size);
    info_simplify(info);
    nms(info);

    // 将检测框重新映射回原始图像尺寸
    for (auto& box : info) 
    {
        box[0] *= (float)img.cols / 640.0;
        box[1] *= (float)img.rows / 640.0;
        box[2] *= (float)img.cols / 640.0;
        box[3] *= (float)img.rows / 640.0;
    }

    return info;
}

void YOLOv5Detector::draw_boxes(Mat& img, const vector<vector<float>>& info)
{
    for (const auto& box : info) 
    {
        if (static_cast<int>(box[5]) == 0) 
        {
            cv::Point topLeft(box[0], box[1]);
            cv::Point bottomRight(box[2], box[3]);
            cv::rectangle(img, topLeft, bottomRight, cv::Scalar(0, 255, 0), 2);
            string label = class_name[0] + " " + std::to_string(box[4]);
            cv::putText(img, label, topLeft, cv::FONT_HERSHEY_SIMPLEX, 0.6, CV_RGB(255, 255, 255), 2);
        }
    }
}

cv::Mat YOLOv5Detector::QImageToMat(const QImage& image)
{
    // 处理 RGB 格式的 QImage
    switch (image.format())
    {
    case QImage::Format_RGB32:
    {
        cv::Mat mat(image.height(), image.width(), CV_8UC4, (void*)image.bits(), image.bytesPerLine());
        cv::Mat bgr;
        cv::cvtColor(mat, bgr, cv::COLOR_RGBA2BGR); // 转换为 BGR
        return bgr.clone(); // 复制以避免修改原始数据
        //return mat.clone(); // 复制以避免修改原始数据
    }
    case QImage::Format_RGB888:
    {
        QImage swapped = image.rgbSwapped(); // OpenCV 使用 BGR 排列
        return cv::Mat(swapped.height(), swapped.width(), CV_8UC3, (void*)swapped.bits(), swapped.bytesPerLine()).clone();
    }
    case QImage::Format_Grayscale8:
    {
        return cv::Mat(image.height(), image.width(), CV_8UC1, (void*)image.bits(), image.bytesPerLine()).clone();
    }
    default:
        break;
    }

    // 如果无法识别格式，返回空的 cv::Mat
    return cv::Mat();
}

vector<vector<float>> YOLOv5Detector::get_info(const float* pdata, int total)
{
    float conf = 0.5;
    int len_data = 85;
    vector<vector<float>> info;
    for (int i = 0; i < total / len_data; i++) 
    {
        if (pdata[4] >= conf) 
        {
            vector<float> line_data(pdata, pdata + len_data);
            info.push_back(line_data);
        }
        pdata += len_data;
    }
    return info;
}

void YOLOv5Detector::info_simplify(vector<vector<float>>& info)
{
    for (int i = 0; i < info.size(); i++) 
    {
        auto max_pos = std::max_element(info[i].cbegin() + 5, info[i].cend());
        int class_id = std::distance(info[i].cbegin() + 5, max_pos);
        info[i][5] = class_id;
        info[i].resize(6);

        float x = info[i][0];
        float y = info[i][1];
        float w = info[i][2];
        float h = info[i][3];
        info[i][0] = x - w / 2.0;
        info[i][1] = y - h / 2.0;
        info[i][2] = x + w / 2.0;
        info[i][3] = y + h / 2.0;
    }
}

void YOLOv5Detector::nms(vector<vector<float>>& info) 
{
    float iou = 0.4;
    int counter = 0;
    vector<vector<float>> return_info;
    while (counter < info.size()) 
    {
        return_info.clear();
        float x1 = 0;
        float x2 = 0;
        float y1 = 0;
        float y2 = 0;

        // 按照置信度排序
        std::sort(info.begin(), info.end(), [](vector<float> p1, vector<float> p2) 
            {
            return p1[4] > p2[4];
            });

        for (auto i = 0; i < info.size(); i++)
        {
            if (i < counter) 
            {
                return_info.push_back(info[i]);
                continue;
            }
            if (i == counter) 
            {
                x1 = info[i][0];
                y1 = info[i][1];
                x2 = info[i][2];
                y2 = info[i][3];
                return_info.push_back(info[i]);
                continue;
            }
            if (info[i][0] > x2 || info[i][2] < x1 || info[i][1] > y2 || info[i][3] < y1)
            {
                return_info.push_back(info[i]);
            }
            else 
            {
                float over_x1 = std::max(x1, info[i][0]);
                float over_y1 = std::max(y1, info[i][1]);
                float over_x2 = std::min(x2, info[i][2]);
                float over_y2 = std::min(y2, info[i][3]);
                float s_over = (over_x2 - over_x1) * (over_y2 - over_y1);
                float s_total = (x2 - x1) * (y2 - y1) + (info[i][2] - info[i][0]) * (info[i][3] - info[i][1]);
                if (s_over / s_total < iou) 
                {
                    return_info.push_back(info[i]);
                }
            }
        }
        info = return_info;
        counter++;
    }
}