#include <filesystem>
#include <vector>
#include <string>
#include <memory>

#include <opencv2/opencv.hpp>

#include "../src/utils_2d_pose.hpp"

// =================================================================================================
// =================================================================================================

int main(int argc, char **argv)
{
    using namespace utils_2d_pose;

    std::string base_path = "/RapidPoseTriangulation/extras/mmdeploy/exports/";
    std::string model_path1 = base_path + "rtmdet-nano_1x320x320x3_fp16_extra-steps.onnx";
    std::string model_path2 = base_path + "rtmpose-m_1x384x288x3_fp16_extra-steps.onnx";

    std::vector<std::string> img_paths = {
        "/RapidPoseTriangulation/data/h1/54138969-img_003201.jpg",
        "/RapidPoseTriangulation/data/h1/55011271-img_003201.jpg",
        "/RapidPoseTriangulation/data/h1/58860488-img_003201.jpg",
        "/RapidPoseTriangulation/data/h1/60457274-img_003201.jpg",
    };

    // RTMDet model1(model_path1, 0.3, 0.1 * 0.1, 30);
    // RTMPose model2(model_path2, 30);
    // for (size_t i = 0; i < img_paths.size(); i++)
    // {
    //     cv::Mat img = cv::imread(img_paths[i]);
    //     cv::cvtColor(img, img, cv::COLOR_BGR2RGB);

    //     auto outputs1 = model1.call(img);
    //     std::cout << "Model1 outputs: " << outputs1[0][0] << " " << outputs1[0][1] << " "
    //               << outputs1[0][2] << " " << outputs1[0][3] << " " << outputs1[0][4] << " "
    //               << std::endl;

    //     for (size_t j = 0; j < outputs1.size(); j++)
    //     {
    //         std::vector<std::array<float, 5>> bboxes = {outputs1[j]};
    //         auto outputs2 = model2.call(img, bboxes);
    //         std::cout << "Model2 outputs: " << outputs2[0][0] << " " << outputs2[0][1] << " "
    //                   << outputs2[0][2] << " " << std::endl;
    //     }
    // }

    // TopDown model3(model_path1, model_path2, 0.3, 0.1 * 0.1, 30);
    // for (size_t i = 0; i < img_paths.size(); i++)
    // {
    //     cv::Mat img = cv::imread(img_paths[i]);
    //     cv::cvtColor(img, img, cv::COLOR_BGR2RGB);

    //     auto outputs3 = model3.predict(img);
    //     std::cout << "Model3 outputs: " << outputs3[0][0][0] << " " << outputs3[0][0][1] << " "
    //               << outputs3[0][0][2] << " " << std::endl;
    // }

    PosePredictor model4(false, 0.3, 0.1 * 0.1, false);
    std::vector<cv::Mat> images;
    for (size_t i = 0; i < img_paths.size(); i++)
    {
        cv::Mat img = cv::imread(img_paths[i]);
        cv::cvtColor(img, img, cv::COLOR_BGR2RGB);
        images.push_back(img);
    }
    auto outputs4 = model4.predict(images);
    std::cout << "Model4 outputs: " << outputs4[0][0][0][0] << " " << outputs4[0][0][0][1] << " "
              << outputs4[0][0][0][2] << " " << std::endl;

    return 0;
}