std::vector<float> blob2vec(const cv::Mat& input_blob)
{

    const auto channels = input_blob.size[1];
    const auto network_width = input_blob.size[2];
    const auto network_height = input_blob.size[3];
    size_t img_byte_size = network_width * network_height * channels * sizeof(float);  // Allocate a buffer to hold all image elements.
    std::vector<float> input_data = std::vector<float>(network_width * network_height * channels);
    std::memcpy(input_data.data(), input_blob.data, img_byte_size);

    std::vector<cv::Mat> chw;
    for (size_t i = 0; i < channels; ++i)
    {
        chw.emplace_back(cv::Mat(cv::Size(network_width, network_height), CV_32FC1, &(input_data[i * network_width * network_height])));
    }
    cv::split(input_blob, chw);

    return input_data;    
}

int main(int argc, char **argv)
{
    const std::string model_path = "/home/oli/yolov8n.onnx";
    const std::string image_path = "dog.jpg";

    YoloVn yolo;
    cv::Mat img = cv::imread(image_path, cv::IMREAD_COLOR);
    cv::Mat input_image = yolo.preprocess_image(img);

    cv::Mat blob;
    cv::dnn::blobFromImage(input_image, blob, 1.0, cv::Size(), cv::Scalar(), false, false);
    auto blob_vec = blob2vec(blob);

    Ort::Env env(ORT_LOGGING_LEVEL_WARNING, "YOLOv8");
    Ort::SessionOptions session_options;
    session_options.SetIntraOpNumThreads(1);
    Ort::Session session(env, model_path.c_str(), session_options);
    auto memory_info = Ort::MemoryInfo::CreateCpu(OrtDeviceAllocator, OrtMemTypeCPU);

    auto input_shape = session.GetInputTypeInfo(0).GetTensorTypeAndShapeInfo().GetShape();
    auto input_tensor = Ort::Value::CreateTensor<float>(
        memory_info,
        blob_vec.data(),
        blob_vec.size(),
        input_shape.data(),
        input_shape.size()
    );

    std::vector<const char*> input_names = {"images"};
    std::vector<const char*> output_names = {"output0"};

    std::vector<Ort::Value> infer_output = session.Run(
        Ort::RunOptions{nullptr},
        input_names.data(),
        &input_tensor,
        1, // input_count
        output_names.data(),
        1 // output_count
    );
    // Postprocess output
    float* outputs_raw = infer_output[0].GetTensorMutableData<float>();
    std::vector<int64_t> output_shape = infer_output[0].GetTensorTypeAndShapeInfo().GetShape();

    // Wrap the raw output in the format expected by postprocess
    std::vector<std::vector<TensorElement>> wrapped_output(1);
    wrapped_output[0].reserve(output_shape[1] * output_shape[2]);
    for (int64_t i = 0; i < output_shape[1] * output_shape[2]; ++i) {
        wrapped_output[0].push_back(outputs_raw[i]);
    }

    std::vector<std::vector<int64_t>> wrapped_shape = {output_shape};

    // Call postprocess
    std::vector<Detection> detections = yolo.postprocess(wrapped_output, wrapped_shape, img.size());

    // Print detections
    std::cout << "Detection results:\n";
    for (const auto& det : detections)
    {
        std::cout
            << "Class: " << det.label
            << ", Confidence: " << det.score
            << ", Bounding Box: (" << det.bbox.x << ", " << det.bbox.y << ", " 
            << det.bbox.width << ", " << det.bbox.height << ")"
            << std::endl;
    }

    return EXIT_SUCCESS;
}