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

#include <opencv2/opencv.hpp>

#include "/RapidPoseTriangulation/scripts/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",
    };

    // {
    //     std::cout << "\nTesting RTMDet and RTMPose" << std::endl;
    //     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][0] << " "
    //                       << outputs2[0][0][1] << " " << outputs2[0][0][2] << " " << std::endl;
    //         }
    //     }
    // }

    // {
    //     std::cout << "\nTesting TopDown" << std::endl;
    //     TopDown model3(model_path1, model_path2, 0.3, 0.1 * 0.1, false, 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;
    //     }
    // }

    {
        std::cout << "\nTesting PosePredictor 1" << 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;
    }

    {
        std::cout << "\nTesting PosePredictor 2" << std::endl;
        PosePredictor model5(false, 0.3, 0.1 * 0.1, true);
        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 outputs5 = model5.predict(images);
        std::cout << "Model5 outputs: " << outputs5[0][0][0][0] << " "
                  << outputs5[0][0][0][1] << " " << outputs5[0][0][0][2] << " " << std::endl;
    }
    return 0;
}
