#include <atomic>
#include <chrono>
#include <cmath>
#include <iostream>
#include <memory>
#include <mutex>
#include <string>
#include <thread>
#include <vector>

// ROS2
#include <rclcpp/rclcpp.hpp>
#include <sensor_msgs/msg/image.hpp>
#include <std_msgs/msg/string.hpp>

// OpenCV / cv_bridge
#include <cv_bridge/cv_bridge.h>
#include <opencv2/opencv.hpp>

// JSON library
#include "/RapidPoseTriangulation/extras/include/nlohmann/json.hpp"
using json = nlohmann::json;

#include "/RapidPoseTriangulation/scripts/utils_2d_pose.hpp"
#include "/RapidPoseTriangulation/scripts/utils_pipeline.hpp"

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

static const std::string cam_id = "camera01";
static const std::string img_input_topic = "/" + cam_id + "/pylon_ros2_camera_node/image_raw";
static const std::string pose_out_topic = "/poses/" + cam_id;

static const float min_bbox_score = 0.4;
static const float min_bbox_area = 0.1 * 0.1;
static const bool batch_poses = true;

static const std::map<std::string, bool> whole_body = {
    {"foots", true},
    {"face", true},
    {"hands", true},
};

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

class Rpt2DWrapperNode : public rclcpp::Node
{
public:
    Rpt2DWrapperNode(const std::string &node_name)
        : Node(node_name)
    {
        this->stop_flag = false;
        this->last_input_image = cv::Mat();
        this->last_input_time = 0.0;

        // QoS
        rclcpp::QoS qos_profile(1);
        qos_profile.reliable();
        qos_profile.keep_last(1);

        // Setup subscriber
        image_sub_ = this->create_subscription<sensor_msgs::msg::Image>(
            img_input_topic, qos_profile,
            std::bind(&Rpt2DWrapperNode::callbackImages, this, std::placeholders::_1));

        // Setup publisher
        pose_pub_ = this->create_publisher<std_msgs::msg::String>(pose_out_topic, qos_profile);

        // Load model
        bool use_wb = utils_pipeline::use_whole_body(whole_body);
        this->kpt_model = std::make_unique<utils_2d_pose::PosePredictor>(
            use_wb, min_bbox_score, min_bbox_area, batch_poses);

        RCLCPP_INFO(this->get_logger(), "Finished initialization of pose estimator.");

        // Start background prediction thread
        model_thread = std::thread(&Rpt2DWrapperNode::callbackWrapper, this);
    }

    ~Rpt2DWrapperNode()
    {
        stop_flag = true;
        if (model_thread.joinable())
        {
            model_thread.join();
        }
    }

private:
    rclcpp::Subscription<sensor_msgs::msg::Image>::SharedPtr image_sub_;
    rclcpp::Publisher<std_msgs::msg::String>::SharedPtr pose_pub_;

    // Pose model pointer
    std::unique_ptr<utils_2d_pose::PosePredictor> kpt_model;
    const std::vector<std::string> joint_names_2d = utils_pipeline::get_joint_names(whole_body);

    // Threading
    std::thread model_thread;
    std::mutex mutex;
    std::atomic<bool> stop_flag;

    cv::Mat last_input_image;
    double last_input_time;

    void callbackImages(const sensor_msgs::msg::Image::SharedPtr msg);
    void callbackModel();

    void callbackWrapper()
    {
        using namespace std::chrono_literals;
        while (!stop_flag)
        {
            callbackModel();
            std::this_thread::sleep_for(std::chrono::microseconds(100));
        }
    }

    void publish(const json &data)
    {
        std_msgs::msg::String msg;
        msg.data = data.dump();
        pose_pub_->publish(msg);
    }
};

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

void Rpt2DWrapperNode::callbackImages(const sensor_msgs::msg::Image::SharedPtr msg)
{
    // Load or convert image to Bayer format
    cv::Mat bayer_image;
    try
    {
        if (msg->encoding == "mono8")
        {
            cv_bridge::CvImageConstPtr cv_ptr = cv_bridge::toCvShare(msg, msg->encoding);
            bayer_image = cv_ptr->image;
        }
        else if (msg->encoding == "bayer_rggb8")
        {
            cv_bridge::CvImageConstPtr cv_ptr = cv_bridge::toCvShare(msg, msg->encoding);
            bayer_image = cv_ptr->image;
        }
        else if (msg->encoding == "rgb8")
        {
            cv_bridge::CvImageConstPtr cv_ptr = cv_bridge::toCvShare(msg, "rgb8");
            cv::Mat color_image = cv_ptr->image;
            bayer_image = utils_pipeline::rgb2bayer(color_image);
        }
        else
        {
            throw std::runtime_error("Unknown image encoding: " + msg->encoding);
        }
    }
    catch (const std::exception &e)
    {
        RCLCPP_ERROR(this->get_logger(), "cv_bridge exception: %s", e.what());
        return;
    }

    // Get timestamp
    double time_stamp = msg->header.stamp.sec + msg->header.stamp.nanosec / 1.0e9;

    // Store in member variables with lock
    {
        std::lock_guard<std::mutex> lk(mutex);
        this->last_input_image = std::move(bayer_image);
        this->last_input_time = time_stamp;
    }
}

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

void Rpt2DWrapperNode::callbackModel()
{
    auto ptime = std::chrono::high_resolution_clock::now();

    // Check if we have an image
    cv::Mat local_image;
    double local_timestamp = 0.0;
    {
        std::lock_guard<std::mutex> lk(mutex);
        if (last_input_time == 0.0)
        {
            return;
        }
        local_image = std::move(last_input_image);
        local_timestamp = last_input_time;

        last_input_image = cv::Mat();
        last_input_time = 0.0;
    }

    // Create image vector
    cv::Mat rgb_image = utils_pipeline::bayer2rgb(local_image);
    std::vector<cv::Mat> images_2d;
    images_2d.push_back(rgb_image);

    // Predict 2D poses
    auto poses_2d_all = kpt_model->predict(images_2d);
    auto poses_2d_upd = utils_pipeline::update_keypoints(
        poses_2d_all, joint_names_2d, whole_body);
    auto &poses_2d = poses_2d_upd[0];

    // Drop persons with no joints
    std::vector<std::vector<std::array<float, 3>>> valid_poses;
    for (auto &person : poses_2d)
    {
        float sum_conf = 0.0;
        for (auto &kp : person)
        {
            sum_conf += kp[2];
        }
        if (sum_conf > 0.0)
        {
            valid_poses.push_back(person);
        }
    }

    // Round poses to 3 decimal places
    for (auto &person : valid_poses)
    {
        for (auto &kp : person)
        {
            kp[0] = std::round(kp[0] * 1000.0) / 1000.0;
            kp[1] = std::round(kp[1] * 1000.0) / 1000.0;
            kp[2] = std::round(kp[2] * 1000.0) / 1000.0;
        }
    }

    // Build JSON
    auto ts_pose = std::chrono::high_resolution_clock::now();
    double ts_pose_sec = std::chrono::duration<double>(ts_pose.time_since_epoch()).count();
    double z_images_pose = ts_pose_sec - local_timestamp;

    json poses_msg;
    poses_msg["bodies2D"] = valid_poses; // shape: persons x keypoints x 3
    poses_msg["joints"] = joint_names_2d;
    poses_msg["num_persons"] = valid_poses.size();
    poses_msg["timestamps"] = {
        {"image", local_timestamp},
        {"pose", ts_pose_sec},
        {"z-images-pose", z_images_pose}};

    // Publish
    publish(poses_msg);

    // Print info
    double elapsed_time = std::chrono::duration<double>(
                              std::chrono::high_resolution_clock::now() - ptime)
                              .count();
    std::cout << "Detected persons: " << valid_poses.size()
              << " - Prediction time: " << elapsed_time << "s" << std::endl;
}

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

int main(int argc, char **argv)
{
    rclcpp::init(argc, argv);

    auto node = std::make_shared<Rpt2DWrapperNode>("rpt2D_wrapper");
    rclcpp::spin(node);

    rclcpp::shutdown();
    return 0;
}