#include "nodes/base/base.hpp"
#include "nodes/infer/inferNode.hpp"
#include "common/image.hpp"
#include "common/utils.hpp"
#include <unordered_map>
#include <random>
#include <algorithm>

namespace GNode
{


void print_mat(const cv::Mat& mat, int max_rows = 10, int max_cols = 10)
{
    for (int i = 0; i < std::min(max_rows, mat.rows); i++) {
        for (int j = 0; j < std::min(max_cols, mat.cols); j++) {
            std::cout << mat.at<float>(i,j) << " ";
        }
        std::cout << (mat.cols > max_cols ? "..." : "") << std::endl;
    }
    if (mat.rows > max_rows) std::cout << "[...]" << std::endl;
}

void InferNode::work()
{
    printf("InferNode %s\n", name_.c_str());
    while (running_)
    {
        Timer timer("InferNode");
        bool has_data = false;
        for (auto& input_buffer : input_buffers_)
        {
            std::shared_ptr<meta::MetaData> metaData;
            if (!input_buffer.second->try_pop(metaData))
            {
                continue;
            }
            has_data = true;
            // printf("Node %s get data from %s\n", name_.c_str(), input_buffer.first.c_str());
            cv::Mat image = metaData->image;
            int width = image.cols;
            int height = image.rows;

            // auto res = model_->forward(tensor::cvimg(image), image.cols, image.rows, 0.0f, 0.0f);
            if (!model_)
            {
                printf("model is nullptr\n");
                continue;
            }
            auto det_result = model_->forward(tensor::cvimg(image), image.cols, image.rows, 0.0f, 0.0f);

            if (std::holds_alternative<data::BoxArray>(det_result)) 
            {
                auto result = std::get<data::BoxArray>(det_result);
                for (auto& r : result)
                {
                    metaData->boxes.push_back(r);
                }
            } 
            else if(std::holds_alternative<cv::Mat>(det_result))
            {
                auto depth_mat = std::get<cv::Mat>(det_result);
                // print_mat(depth_mat);
                metaData->depth = depth_mat;

            }
            else
            {
                printf("Unexpected result type from model");
                throw std::runtime_error("Unexpected result type from model");
            }
            for (auto& output_buffer : output_buffers_)
            {
                // printf("Node %s push data to %s\n", name_.c_str(), output_buffer.first.c_str());
                output_buffer.second->push(metaData);
            }           
        }
        if (!has_data)
        {
            std::unique_lock<std::mutex> lock(mutex_);
            cond_var_->wait_for(lock, std::chrono::milliseconds(100), [this] {
                return !running_;  // 等待时检查退出条件
            });
        }
    }
};

}   // namespace StreamNode