#include <ros/ros.h>
#include <tf/transform_broadcaster.h>
#include <geometry_msgs/Pose.h>
#include <chrono>
#include <comon/clockhelper.hpp>
#include <pcl_ros/point_cloud.h>
#include <pcl/point_cloud.h>
#include <pcl/point_types.h>
#include <pcl/filters/voxel_grid.h>
#include <pcl/io/pcd_io.h>
#include <pcl/io/ply_io.h>
#include <pcl/point_types.h>
#include <pcl/filters/voxel_grid.h>
#include <pcl/filters/statistical_outlier_removal.h>
#include <pcl/filters/passthrough.h>
#include <pcl/registration/icp.h>
#include <pcl/registration/ia_ransac.h>
#include <pcl/features/pfh.h>
#include <pcl/features/fpfh.h>
#include <pcl/features/normal_3d.h>
#include <Eigen/Geometry>
#include <pcl/common/transforms.h>
#include <comon/rosbaghelper.hpp>
#include <comon/pclhelper.hpp>


// Less typing
using CloudRGBA = pcl::PointCloud<pcl::PointXYZRGBA>;

void setCloudColor(CloudRGBA::Ptr cloud, int r, int g, int b) {
    for (int i = 0; i < cloud->size(); ++i) {
        auto& point = (*cloud)[i];
        point.r = r;
        /* MISSING - Task 2 */
        /* MISSING - Task 2 */
    }
}

void prepareHeader(CloudRGBA::Ptr cloud, const std::string frameID) {
    pcl_conversions::toPCL(ros::Time::now(), cloud->header.stamp);
    cloud->header.frame_id = ""; /* FIX ME - Task 2 */
}

CloudRGBA::Ptr downsample(CloudRGBA::Ptr frame, float leafSize) {
    auto downsampledCloud = boost::make_shared<CloudRGBA>();
    if (leafSize == 0) {
        *downsampledCloud = *frame;
        return downsampledCloud;
    }
    // Downsample using voxelgrid
    pcl::VoxelGrid<pcl::PointXYZRGBA> downsample_filter;
    downsample_filter.setLeafSize(0, 0, 0); /* FIX ME - Task 2 */
    downsample_filter.setInputCloud(frame);
    downsample_filter.filter(*downsampledCloud);

    return downsampledCloud;
}

int main(int argc, char **argv) {
    // The name of this node
    std::string nodeName = "crs_tf_publisher";

    ros::init(argc, argv, nodeName);
    auto node = ros::NodeHandle();

    // The used bag file
    std::string bagFilePath = "/home/marius/bockdata/bags/room2.bag";


    ros::Publisher framePublisher, modelPublisher, alignedPublisher, completeCloudPublisher;
    framePublisher = node.advertise<CloudRGBA>("/crs/frame", 5);
    modelPublisher = node.advertise<CloudRGBA>("", 5); /* FIX ME - Task 3 */
    alignedPublisher = node.advertise<CloudRGBA>("", 5); /* FIX ME - Task 3 */
    completeCloudPublisher = node.advertise<CloudRGBA>("", 5); /* FIX ME - Task 3 */

    // The topic we subscribe to
    std::string defaultCloudFrame = "kinect2_ir_optical_frame";

    // Some ICP Parameters. DO NOT CHANGE THEM!
    double ICPMaxCorrespondenceDistance = 0.2;
    double ICPTransformationEpsilon = 0.00001;
    double ICPEuclidianFitnessEpsilon = 0.1;


    // The leaf size of the VoxelGrid Filter
    double ICPDownsampleLeafSize = 0.1;

    // The maximal of ICP iterations
    int ICPMaxIterations = 20;

    // The ROS Bag Helper.
    BagFileHelper helper;
    helper.openBagFile(bagFilePath);
    helper.initTopic(cloudSubscribeTopic);

    // Create an empty cloud that is going to contain ALL registered point clouds
    auto completeCloud = boost::make_shared<CloudRGBA>();

    // Create a matrix that is going to contain the transformation from the CURRENT frame to the FIRST Model
    // At the beginning it is an identiy matrix -> no transformation at all
    Eigen::Matrix4f completeTransform;
    completeTransform.setIdentity();

    // Variables for the model and the frame
    CloudRGBA::Ptr model, frame;

    // The main loop runs as long as there are new clouds avaiable and [CRTL + c] has not been pressed
    while (helper.hasTopicAvailableMessages("/kinect2/sd/points") && ros::ok()) {
        // Gather the next cloud from our bag file
        frame = helper.getNextCloudRGBAForTopic("/kinect2/sd/points");

        if (model == nullptr) {
            // Since it is the first loop iteration we do not have a model yet.
            // The current frame will be the model for the next frame
            model = frame;
        } else {
            // Downsample the input clouds
            auto downsampledModel = downsample(model, ICPDownsampleLeafSize);
            auto downsampledFrame = downsample(frame, ICPDownsampleLeafSize);

            // Create the ICP object
            IterativeClosestPoint icp;

            // Set the ICP parameters
            icp.setMaxCorrespondenceDistance(ICPMaxCorrespondenceDistance);
            icp.setMaximumIterations(ICPMaxIterations);
            icp.setTransformationEpsilon(ICPTransformationEpsilon);
            icp.setEuclideanFitnessEpsilon(ICPEuclidianFitnessEpsilon);

            // Set ICP data sources
            icp.setInputCloud(downsampledFrame);
            icp.setInputTarget(downsampledModel);

            // Run ICP Alligment
            icp.align(*downsampledFrame);

            // The ICP algorithm has finished. Print the convergence criteria!
            ROS_INFO("ICP terminated with criteria: %s", getCriteria(icp));

            // Receive the local transformation from the current frame to the model (the last frame)
            auto transform = icp.getFinalTransformation();

            // Calculate the complete transformation from the CURRENT frame to the FIRST model
            completeTransform = transform * completeTransform;

            // Transform the current frame and add it to the complete model
            auto transformedCloud = boost::make_shared<CloudRGBA>();
            pcl::transformPointCloud(*frame, *transformedCloud, completeTransform);
            *completeCloud += *transformedCloud;

            // Publish the complete model
            prepareHeader(completeCloud, defaultCloudFrame);
            completeCloudPublisher.publish(completeCloud);

            // Publish the current model
            setCloudColor(downsampledModel, 0, 0, 0); /* FIX ME - Task 3 */
            prepareHeader(downsampledModel, defaultCloudFrame);
            modelPublisher.publish(downsampledModel);

            // Publish the current frame
            setCloudColor(downsampledFrame, 0, 0, 0); /* FIX ME - Task 3 */
            prepareHeader(downsampledFrame, defaultCloudFrame);
            framePublisher.publish(downsampledFrame);

            // Publish the current, donwsampled aligned frame
            auto downsampledAlign = downsample(transformedCloud, ICPDownsampleLeafSize);
            setCloudColor(downsampledAlign, 0, 0, 0); /* FIX ME - Task 3 */
            prepareHeader(downsampledAlign, defaultCloudFrame);
            alignedPublisher.publish(downsampledAlign);

            // The current frame will be the model for the next frame
            model = frame;
        }
    }
    return 0;
}