#include <atomic>
#include <chrono>
#include <iostream>
#include <mutex>
#include <string>
#include <vector>

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

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

#include "rpt_msgs/msg/poses.hpp"
#include "/RapidPoseTriangulation/rpt/camera.hpp"
#include "/RapidPoseTriangulation/rpt/interface.hpp"
#include "/RapidPoseTriangulation/rpt/tracker.hpp"

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

static const std::vector<std::string> cam_ids = {
    "camera01",
    "camera02",
    "camera03",
    "camera04",
    "camera05",
    "camera06",
    "camera07",
    "camera08",
    "camera09",
    "camera10",
};

static const std::string pose_in_topic = "/poses/{}";
static const std::string cam_info_topic = "/{}/calibration";
static const std::string pose_out_topic = "/poses/humans3d";

static const float min_match_score = 0.94;
static const size_t min_group_size = 1;

static const bool use_tracking = true;
static const float max_movement_speed = 2.0 * 1.5;
static const float cam_fps = 50;
static const float max_track_distance = 0.3 + max_movement_speed / cam_fps;

static const std::array<std::array<float, 3>, 2> roomparams = {{
    {4.0, 5.0, 2.2},
    {1.0, 0.0, 1.1},
}};

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

class Rpt3DWrapperNode : public rclcpp::Node
{
public:
    Rpt3DWrapperNode()
        : Node("rpt3d_wrapper")
    {
        this->all_cameras.resize(cam_ids.size());
        this->all_poses.resize(cam_ids.size());
        this->all_timestamps.resize(cam_ids.size());
        this->joint_names = {};
        this->all_poses_set.resize(cam_ids.size(), false);

        // Load 3D models
        tri_model = std::make_unique<Triangulator>(
            min_match_score, min_group_size);
        pose_tracker = std::make_unique<PoseTracker>(
            max_movement_speed, max_track_distance);

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

        // Parallel executable callbacks
        auto my_callback_group = create_callback_group(rclcpp::CallbackGroupType::Reentrant);
        rclcpp::SubscriptionOptions options;
        options.callback_group = my_callback_group;

        // Setup subscribers
        for (size_t i = 0; i < cam_ids.size(); i++)
        {
            std::string cam_id = cam_ids[i];
            std::string topic_pose = pose_in_topic;
            std::string topic_cam = cam_info_topic;
            topic_pose.replace(topic_pose.find("{}"), 2, cam_id);
            topic_cam.replace(topic_cam.find("{}"), 2, cam_id);

            auto sub_pose = this->create_subscription<rpt_msgs::msg::Poses>(
                topic_pose, qos_profile,
                [this, i](const rpt_msgs::msg::Poses::SharedPtr msg)
                {
                    this->callback_poses(msg, i);
                },
                options);
            sub_pose_list_.push_back(sub_pose);

            auto sub_cam = this->create_subscription<std_msgs::msg::String>(
                topic_cam, qos_profile,
                [this, i](const std_msgs::msg::String::SharedPtr msg)
                {
                    this->callback_cam_info(msg, i);
                },
                options);
            sub_cam_list_.push_back(sub_cam);
        }

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

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

private:
    std::vector<rclcpp::Subscription<rpt_msgs::msg::Poses>::SharedPtr> sub_pose_list_;
    std::vector<rclcpp::Subscription<std_msgs::msg::String>::SharedPtr> sub_cam_list_;
    rclcpp::Publisher<rpt_msgs::msg::Poses>::SharedPtr pose_pub_;

    std::unique_ptr<Triangulator> tri_model;
    std::unique_ptr<PoseTracker> pose_tracker;
    std::vector<Camera> all_cameras;
    std::mutex cams_mutex, pose_mutex, model_mutex;

    std::vector<std::vector<std::vector<std::array<float, 3>>>> all_poses;
    std::vector<double> all_timestamps;
    std::vector<std::string> joint_names;
    std::vector<bool> all_poses_set;

    void callback_poses(const rpt_msgs::msg::Poses::SharedPtr msg, size_t cam_idx);
    void callback_cam_info(const std_msgs::msg::String::SharedPtr msg, size_t cam_idx);
    void call_model();
};

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

void Rpt3DWrapperNode::callback_cam_info(const std_msgs::msg::String::SharedPtr msg, size_t cam_idx)
{
    json cam = json::parse(msg->data);

    Camera camera;
    camera.name = cam["name"].get<std::string>();
    camera.K = cam["K"].get<std::array<std::array<float, 3>, 3>>();
    camera.DC = cam["DC"].get<std::vector<float>>();
    camera.R = cam["R"].get<std::array<std::array<float, 3>, 3>>();
    camera.T = cam["T"].get<std::array<std::array<float, 1>, 3>>();
    camera.width = cam["width"].get<int>();
    camera.height = cam["height"].get<int>();
    camera.type = cam["type"].get<std::string>();

    cams_mutex.lock();
    all_cameras[cam_idx] = camera;
    cams_mutex.unlock();
}

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

void Rpt3DWrapperNode::callback_poses(const rpt_msgs::msg::Poses::SharedPtr msg, size_t cam_idx)
{
    std::vector<std::vector<std::array<float, 3>>> poses;
    if (this->joint_names.empty())
    {
        this->joint_names = msg->joint_names;
    }

    // Unflatten poses
    size_t idx = 0;
    std::vector<int32_t> &bshape = msg->bodies_shape;
    for (int32_t i = 0; i < bshape[0]; ++i)
    {
        std::vector<std::array<float, 3>> body;
        for (int32_t j = 0; j < bshape[1]; ++j)
        {
            std::array<float, 3> joint;
            for (int32_t k = 0; k < bshape[2]; ++k)
            {
                joint[k] = msg->bodies_flat[idx];
                ++idx;
            }
            body.push_back(std::move(joint));
        }
        poses.push_back(std::move(body));
    }

    // If no pose was detected, create an empty placeholder
    if (poses.size() == 0)
    {
        std::vector<std::array<float, 3>> body(joint_names.size(), {0, 0, 0});
        poses.push_back(std::move(body));
    }

    pose_mutex.lock();
    this->all_poses[cam_idx] = std::move(poses);
    this->all_poses_set[cam_idx] = true;
    this->all_timestamps[cam_idx] = msg->header.stamp.sec + msg->header.stamp.nanosec / 1.0e9;
    pose_mutex.unlock();

    // Trigger model callback
    model_mutex.lock();
    call_model();
    model_mutex.unlock();
}

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

void Rpt3DWrapperNode::call_model()
{
    auto ts_msg = std::chrono::high_resolution_clock::now();

    // Check if all cameras are set
    cams_mutex.lock();
    for (size_t i = 0; i < cam_ids.size(); i++)
    {
        if (all_cameras[i].name.empty())
        {
            RCLCPP_WARN(this->get_logger(), "Skipping frame, still waiting for cameras...");
            cams_mutex.unlock();
            return;
        }
    }
    cams_mutex.unlock();

    // If not all poses are set, return and wait for the others
    pose_mutex.lock();
    for (size_t i = 0; i < cam_ids.size(); i++)
    {
        if (!this->all_poses_set[i])
        {
            pose_mutex.unlock();
            return;
        }
    }
    pose_mutex.unlock();

    // Call model, and meanwhile lock updates of the inputs
    // Since the prediction is very fast, parallel callback threads only need to wait a short time
    cams_mutex.lock();
    pose_mutex.lock();
    const auto poses_3d = tri_model->triangulate_poses(
        all_poses, all_cameras, roomparams, joint_names);

    double min_ts = *std::min_element(all_timestamps.begin(), all_timestamps.end());
    this->all_poses_set = std::vector<bool>(cam_ids.size(), false);

    std::vector<std::vector<std::array<float, 4>>> valid_poses;
    std::vector<size_t> track_ids;
    if (use_tracking)
    {
        auto pose_tracks = pose_tracker->track_poses(poses_3d, joint_names, min_ts);
        std::vector<std::vector<std::array<float, 4>>> poses_3d_refined;
    for (size_t j = 0; j < pose_tracks.size(); j++)
    {
        auto &pose = std::get<1>(pose_tracks[j]);
        poses_3d_refined.push_back(pose);
        auto &track_id = std::get<0>(pose_tracks[j]);
        track_ids.push_back(track_id);
    }
        valid_poses = std::move(poses_3d_refined);
    }
    else
    {
        valid_poses = std::move(poses_3d);
        track_ids = {};
    }
    pose_mutex.unlock();
    cams_mutex.unlock();

    // Calculate timings
    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_trigger_pose3d = ts_pose_sec - min_ts;
    json jdata;
    jdata["timestamps"] = {
        {"trigger", min_ts},
        {"pose3d", ts_pose_sec},
        {"z-trigger-pose3d", z_trigger_pose3d}};

    // Publish message
    auto pose_msg = rpt_msgs::msg::Poses();
    pose_msg.header.stamp.sec = static_cast<int>(min_ts);
    pose_msg.header.stamp.nanosec = (min_ts - pose_msg.header.stamp.sec) * 1.0e9;
    pose_msg.header.frame_id = "world";
    std::vector<int32_t> pshape = {(int)valid_poses.size(), (int)joint_names.size(), 4};
    pose_msg.bodies_shape = pshape;
    pose_msg.bodies_flat.reserve(pshape[0] * pshape[1] * pshape[2]);
    for (int32_t i = 0; i < pshape[0]; i++)
    {
        for (int32_t j = 0; j < pshape[1]; j++)
        {
            for (int32_t k = 0; k < pshape[2]; k++)
            {
                pose_msg.bodies_flat.push_back(valid_poses[i][j][k]);
            }
        }
    }
    pose_msg.joint_names = joint_names;
    jdata["track_ids"] = track_ids;
    pose_msg.extra_data = jdata.dump();
    pose_pub_->publish(pose_msg);

    // Print info
    double elapsed_time = std::chrono::duration<double>(
                              std::chrono::high_resolution_clock::now() - ts_msg)
                              .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<Rpt3DWrapperNode>();
    rclcpp::executors::MultiThreadedExecutor exec;
    exec.add_node(node);
    exec.spin();
    rclcpp::shutdown();
    return 0;
}
