/**
  * @project Project
  * @brief   TODO: add a brief
  * @author  lz
  * @data    2023/6/24
 **/

#include <iostream>
#include <string>
#include <pcl/io/ply_io.h>
#include <pcl/point_types.h>
#include <pcl/registration/icp.h>
#include <pcl/visualization/pcl_visualizer.h>
#include <pcl/console/time.h>
#include <vector>
#include <omp.h>

#include <map>
#include "rslidar_mqtt_async.h"

using namespace std;

typedef pcl::PointXYZ PointT;
typedef pcl::PointCloud<PointT> PointCloudT;
bool next_iteration = false;

void print4x4Matrix(const Eigen::Matrix4d &matrix) {
    printf("Rotation matrix :\n");
    printf("    | %6.6f %6.6f %6.6f | \n", matrix(0, 0), matrix(0, 1), matrix(0, 2));
    printf("R = | %6.6f %6.6f %6.6f | \n", matrix(1, 0), matrix(1, 1), matrix(1, 2));
    printf("    | %6.6f %6.6f %6.6f | \n", matrix(2, 0), matrix(2, 1), matrix(2, 2));
    printf("Translation vector :\n");
    printf("t = < %6.6f, %6.6f, %6.6f >\n\n", matrix(0, 3), matrix(1, 3), matrix(2, 3));
}

void keyboardEventOccurred(const pcl::visualization::KeyboardEvent &event, void *nothing) {
    if (event.getKeySym() == "space" && event.keyDown())
        next_iteration = true;
}

void findcorrespondences(pcl::PointCloud<pcl::PointXYZ>::Ptr &source, pcl::PointCloud<pcl::PointXYZ>::Ptr &target,
                         pcl::PointCloud<pcl::PointXYZ>::Ptr &P, pcl::PointCloud<pcl::PointXYZ>::Ptr &Q) {
    P->clear();
    Q->clear();
    pcl::KdTreeFLANN<pcl::PointXYZ> kdtree;
    kdtree.setInputCloud(target);
//#pragma omp parallel for num_threads(200)
    for (int i = 0; i < source->points.size(); i++) {
        vector<int> indx;
        vector<float> dists;
        kdtree.nearestKSearch(source->points[i], 1, indx, dists);
        if (dists[0] < 0.02) {
            P->push_back(source->points[i]);
            Q->push_back(target->points[indx[0]]);
        }
    }
}

Eigen::Matrix<double, 4, 4> findRT(pcl::PointCloud<pcl::PointXYZ>::Ptr &P, pcl::PointCloud<pcl::PointXYZ>::Ptr &Q) {
    int n = P->points.size();
    Eigen::Matrix<double, 3, Eigen::Dynamic> matrixP(3, n), matrixQ(3, n);
    for (int i = 0; i < n; i++) {
        matrixP(0, i) = P->points[i].x;
        matrixP(1, i) = P->points[i].y;
        matrixP(2, i) = P->points[i].z;
        matrixQ(0, i) = Q->points[i].x;
        matrixQ(1, i) = Q->points[i].y;
        matrixQ(2, i) = Q->points[i].z;
    }
    Eigen::Matrix4d RT = Eigen::umeyama(matrixP, matrixQ);
    print4x4Matrix(RT);
    return RT;
}

void myICP(pcl::PointCloud<pcl::PointXYZ>::Ptr &source, pcl::PointCloud<pcl::PointXYZ>::Ptr &target,
           Eigen::Matrix<double, 4, 4> &RT) {
    pcl::PointCloud<pcl::PointXYZ>::Ptr P(new pcl::PointCloud<pcl::PointXYZ>);
    pcl::PointCloud<pcl::PointXYZ>::Ptr Q(new pcl::PointCloud<pcl::PointXYZ>);

    pcl::visualization::PCLVisualizer viewer("viewer");

    pcl::visualization::PointCloudColorHandlerCustom<PointT> cloud_source(source, 255, 0, 0), cloud_target(target, 255,
                                                                                                           255, 255);

    viewer.addPointCloud(source, cloud_source, "source");
    viewer.addPointCloud(target, cloud_target, "target");

    viewer.registerKeyboardCallback(&keyboardEventOccurred, (void *) nullptr);
    while (!viewer.wasStopped()) {
        viewer.spinOnce();
        if (next_iteration) {
            findcorrespondences(source, target, P, Q);
            RT = findRT(P, Q);
            pcl::transformPointCloud(*source, *source, RT);
            viewer.updatePointCloud(source, cloud_source, "source");
        }
        next_iteration = false;
    }
}


int main() {
    CloudDataMqttAsync mqtt_client;
    mqtt_client.init();

    std::map<std::string, pcl::PointCloud<pcl::PointXYZ>::Ptr> cloud_map;
    cloud_map.insert(std::pair<std::string, pcl::PointCloud<pcl::PointXYZ>::Ptr>("rslidar/6311",
                                                                                 new pcl::PointCloud<pcl::PointXYZ>));
    cloud_map.insert(std::pair<std::string, pcl::PointCloud<pcl::PointXYZ>::Ptr>("rslidar/6312",
                                                                                 new pcl::PointCloud<pcl::PointXYZ>));
    cloud_map.insert(std::pair<std::string, pcl::PointCloud<pcl::PointXYZ>::Ptr>("rslidar/6313",
                                                                                 new pcl::PointCloud<pcl::PointXYZ>));
    cloud_map.insert(std::pair<std::string, pcl::PointCloud<pcl::PointXYZ>::Ptr>("rslidar/6314",
                                                                                 new pcl::PointCloud<pcl::PointXYZ>));
    cloud_map.insert(std::pair<std::string, pcl::PointCloud<pcl::PointXYZ>::Ptr>("rslidar/6315",
                                                                                 new pcl::PointCloud<pcl::PointXYZ>));
    cloud_map.insert(std::pair<std::string, pcl::PointCloud<pcl::PointXYZ>::Ptr>("rslidar/6316",
                                                                                 new pcl::PointCloud<pcl::PointXYZ>));

    std::map<std::string, Eigen::Matrix<double, 4, 4>> rt_map;
    rt_map.insert(std::pair<std::string, Eigen::Matrix<double, 4, 4>>("rslidar/6311", Eigen::Matrix<double, 4, 4>()));
    rt_map.insert(std::pair<std::string, Eigen::Matrix<double, 4, 4>>("rslidar/6312", Eigen::Matrix<double, 4, 4>()));
    rt_map.insert(std::pair<std::string, Eigen::Matrix<double, 4, 4>>("rslidar/6313", Eigen::Matrix<double, 4, 4>()));
    rt_map.insert(std::pair<std::string, Eigen::Matrix<double, 4, 4>>("rslidar/6314", Eigen::Matrix<double, 4, 4>()));
    rt_map.insert(std::pair<std::string, Eigen::Matrix<double, 4, 4>>("rslidar/6315", Eigen::Matrix<double, 4, 4>()));
    rt_map.insert(std::pair<std::string, Eigen::Matrix<double, 4, 4>>("rslidar/6316", Eigen::Matrix<double, 4, 4>()));

    bool run = true;
    while (run) {
        for (auto &iter: cloud_map) {
            mqtt_client.getCloudData(iter.first, iter.second);
        }

        if (cloud_map.find("rslidar/6314")->second->empty()) {
            continue;
        }

        if (!cloud_map.find("rslidar/6311")->second->empty()) {
            myICP(cloud_map.find("rslidar/6314")->second, cloud_map.find("rslidar/6311")->second,
                  rt_map.find("rslidar/6311")->second);
        }

        break;

        if (!cloud_map.find("rslidar/6312")->second->empty()) {
            myICP(cloud_map.find("rslidar/6314")->second, cloud_map.find("rslidar/6312")->second,
                  rt_map.find("rslidar/6312")->second);

        }
        if (!cloud_map.find("rslidar/6313")->second->empty()) {
            myICP(cloud_map.find("rslidar/6314")->second, cloud_map.find("rslidar/6313")->second,
                  rt_map.find("rslidar/6313")->second);
        }

        if (!cloud_map.find("rslidar/6314")->second->empty()) {
            myICP(cloud_map.find("rslidar/6314")->second, cloud_map.find("rslidar/6314")->second,
                  rt_map.find("rslidar/6314")->second);
        }

        if (!cloud_map.find("rslidar/6315")->second->empty()) {
            myICP(cloud_map.find("rslidar/6314")->second, cloud_map.find("rslidar/6315")->second,
                  rt_map.find("rslidar/6315")->second);
        }

        if (!cloud_map.find("rslidar/6316")->second->empty()) {
            myICP(cloud_map.find("rslidar/6314")->second, cloud_map.find("rslidar/6316")->second,
                  rt_map.find("rslidar/6316")->second);
        }

        run = false;
    }
}