#include <iostream>
#include <Eigen/Dense>

#include <pcl/io/pcd_io.h>

#include <pcl/point_cloud.h>
#include <pcl/point_types.h>

#include <pcl/filters/filter.h>

#include <pcl/registration/icp.h>

#include <pcl/common/time.h>

#include <pcl/visualization/pcl_visualizer.h>

bool next_iteration = false;

// 此函数是可视化窗口的回调函数。当可视化窗口被激活时，只要按空格键，就会调用此函数，将布尔变量next_iteration 设置为 true
void keyboardEventOccurred (const pcl::visualization::KeyboardEvent& event, void* nothing)
{
    if (event.getKeySym () == "space" && event.keyDown ())
        next_iteration = true;
}

void down_sample(const pcl::PointCloud<pcl::PointXYZRGB>::Ptr& cloud_in,
                 const pcl::PointCloud<pcl::PointXYZRGB>::Ptr& cloud_out)
{
    pcl::ScopeTime time("down sample");

    int down_sample_rate = 8;

    int height = cloud_in->height;
    int width = cloud_in->width;

    int height_ = height / down_sample_rate;
    int width_ = width / down_sample_rate;

    cloud_out->resize(height_ * width_);

    for (int i = 0; i < height; i = i + down_sample_rate)
    {
        for (int j = 0; j < width; j = j + down_sample_rate)
        {
            int index = i * width + j;
            int index_ = i * width / down_sample_rate / down_sample_rate + j / down_sample_rate;

            cloud_out->points[index_] = cloud_in->points[index];
        }
    }

    cloud_out->height = height_;
    cloud_out->width = width_;
}

void remove_nan(const pcl::PointCloud<pcl::PointXYZRGB>::Ptr& cloud_1,
                const pcl::PointCloud<pcl::PointXYZRGB>::Ptr& cloud_2)
{
    // 如果想正确使用 pcl::removeNaNFromPointCloud 函数，最好先把点云的 is_dense 属性设置为 false
    // 因为 pcl::removeNaNFromPointCloud 函数会先检查 is_dense 属性，如果该属性原本为 true 那么该函数将不会对点云进行任何处理
    cloud_1->is_dense = false;
    cloud_2->is_dense = false;

    std::vector<int> indices_1;
    pcl::removeNaNFromPointCloud(*cloud_1, *cloud_1, indices_1);

    std::vector<int> indices_2;
    pcl::removeNaNFromPointCloud(*cloud_2, *cloud_2, indices_2);

    cloud_1->is_dense = true;
    cloud_2->is_dense = true;
}

void icp_function(const pcl::PointCloud<pcl::PointXYZRGB>::Ptr& cloud_source,
                  const pcl::PointCloud<pcl::PointXYZRGB>::Ptr& cloud_target)
{
    int iterations = 10;

    // ICP 算法
    pcl::IterativeClosestPoint<pcl::PointXYZRGB, pcl::PointXYZRGB> icp;
    icp.setMaximumIterations (iterations);  // 设置要执行的初始迭代次数（默认值为1）
    icp.setInputSource(cloud_source);
    icp.setInputTarget(cloud_target);
    icp.align(*cloud_source);
    icp.setMaximumIterations(1);  // 在第一次对齐之后，将这个变量设置为1，下次调用.align（）函数的时候就只会迭代一次

    if (icp.hasConverged ())
    {
        std::cout << "\nICP has converged, score is " << icp.getFitnessScore () << std::endl;
        std::cout << "\nICP transformation " << iterations << " : cloud_source -> cloud_target" << std::endl;

        Eigen::Matrix4d transformation_matrix = icp.getFinalTransformation ().cast<double>();
        std::cout << transformation_matrix << std::endl;
    }
    else
    {
        std::cerr << "\nICP has not converged." << std::endl;
    }

    // 可视化部分
    pcl::visualization::PCLVisualizer viewer;

    // 目标点云为白色
    pcl::visualization::PointCloudColorHandlerCustom<pcl::PointXYZRGB> white(cloud_target, 250, 250, 250);

    viewer.addPointCloud<pcl::PointXYZRGB> (cloud_target, white, "cloud_target");

    // 原点云为红色
    pcl::visualization::PointCloudColorHandlerCustom<pcl::PointXYZRGB> red (cloud_source, 250, 0, 0);
    viewer.addPointCloud (cloud_source, red, "cloud_source");

    // 键盘回调函数
    viewer.registerKeyboardCallback (&keyboardEventOccurred, (void*) nullptr);

    // 最终可视化的时候会看到，每按一次空格键，红色的原点云会逐渐逼近白色的目标点云
    while (!viewer.wasStopped())
    {
        viewer.spinOnce ();

        // if you pressed "space"
        if (next_iteration)
        {
            icp.align (*cloud_source);  // 每一次循环，只迭代一次

            if (icp.hasConverged ())
            {
                std::cout << "\nICP has converged, score is " << icp.getFitnessScore () << std::endl;
                std::cout << "\nICP transformation " << ++iterations << " : cloud_source -> cloud_target" << std::endl;

                Eigen::Matrix4d transformation_matrix = icp.getFinalTransformation ().cast<double>();
                std::cout << transformation_matrix << std::endl;

                viewer.updatePointCloud (cloud_source, red, "cloud_source");
            }
            else
            {
                std::cerr << "\nICP has not converged.\n" << std::endl;
            }
        }

        next_iteration = false;
    }
}

int main()
{
    pcl::PointCloud<pcl::PointXYZRGB>::Ptr cloud_source (new pcl::PointCloud<pcl::PointXYZRGB>);  // 原点云
    pcl::PointCloud<pcl::PointXYZRGB>::Ptr cloud_target (new pcl::PointCloud<pcl::PointXYZRGB>);  // 目标点云

    std::string name_1 = "XXX.pcd" ;
    pcl::io::loadPCDFile<pcl::PointXYZRGB>(name_1, *cloud_source);

    std::string name_2 = "YYY.pcd" ;
    pcl::io::loadPCDFile<pcl::PointXYZRGB>(name_2, *cloud_target);

    std::cout << "[Cloud 1] Number of Points = " << cloud_source->points.size () << std::endl;
    std::cout << "[Cloud 2] Number of Points = " << cloud_target->points.size () << std::endl;

    // 进行降采样，把原始点云的尺寸从 640 * 480 降采样到 80 * 60
    down_sample(cloud_source, cloud_source);
    down_sample(cloud_target, cloud_target);

    std::cout << "[Cloud 1] Number of Points = " << cloud_source->points.size () << std::endl;
    std::cout << "[Cloud 2] Number of Points = " << cloud_target->points.size () << std::endl;

    // 点云中有些点是无效的，要将其去除，要不然 ICP 会报错
    remove_nan(cloud_source, cloud_target);

    icp_function(cloud_source, cloud_target);

    return 0;
}