#include <librealsense2/rs.hpp>
#include <iostream>
#include <filesystem>
#include <pcl/io/ply_io.h>
#include <pcl/point_types.h>
#include "pcl-helper.hpp"

struct ColoredPoints
{
    rs2::points points;
    rs2::video_frame color;

    ColoredPoints(rs2::points p, rs2::video_frame c) : points(p), color(c) {}
};

void configure_filters(rs2::decimation_filter &dec_filter, rs2::threshold_filter &thr_filter, rs2::spatial_filter &spat_filter, int dec_filter_magnitude, float min_dist, float max_dist, int magnitude, int smooth_alpha, int smooth_delta, int holes_fill);
void apply_filters(rs2::decimation_filter &dec_filter, rs2::threshold_filter &thr_filter, rs2::spatial_filter &spat_filter, rs2::disparity_transform &depth_to_disparity, rs2::disparity_transform &disparity_to_depth, rs2::depth_frame &depth_frame);
ColoredPoints scan(rs2::decimation_filter &dec_filter, rs2::threshold_filter &thr_filter, rs2::spatial_filter &spat_filter, rs2::disparity_transform &depth_to_disparity, rs2::disparity_transform &disparity_to_depth, rs2::pipeline &pipe);
void scan_and_convert(rs2::decimation_filter &dec_filter, rs2::threshold_filter &thr_filter, rs2::spatial_filter &spat_filter, rs2::disparity_transform &depth_to_disparity, rs2::disparity_transform &disparity_to_depth, rs2::pipeline &pipe, std::vector<PCLHelper::PtrCloud> &pointclouds);
bool align_and_save_pointclouds(std::vector<PCLHelper::PtrCloud> &pointclouds);
void save_original_pointclouds(std::vector<PCLHelper::PtrCloud> &pointclouds);
void wipe_old_scans();
int prompt_user();

int main()
try
{
    rs2::pipeline pipe;

    pipe.start();

    // define filters
    rs2::decimation_filter dec_filter; // Decimation - reduces depth frame density
    rs2::threshold_filter thr_filter;  // Threshold  - removes values outside recommended range
    rs2::spatial_filter spat_filter;   // Spatial    - edge-preserving spatial smoothing
    // rs2::hole_filling_filter hole_filter;                    // Hole filling - fills small holes in the depth image
    const std::string disparity_filter_name = "Disparity"; // Declare disparity transform from depth to disparity and vice versa
    rs2::disparity_transform depth_to_disparity(true);
    rs2::disparity_transform disparity_to_depth(false);
    configure_filters(dec_filter, thr_filter, spat_filter, 3, 0.0, 0.4, 5, 1, 1, 5);

    std::vector<PCLHelper::PtrCloud> pointclouds;
    PCLHelper::PtrCloud converted_points;
    PCLHelper::PtrCloud cropped_points;
    bool running = true;
    while (running)
    {

        int number = prompt_user();

        switch (number)
        {
        case 1:
            scan_and_convert(dec_filter, thr_filter, spat_filter, depth_to_disparity, disparity_to_depth, pipe, pointclouds);
            break;
        case 2:
            if (pointclouds.size() < 2)
            {
                std::cout << "Not enough pointclouds to merge, just exiting...";
            }
            else
            {
                save_original_pointclouds(pointclouds);
                if (!align_and_save_pointclouds(pointclouds))
                {
                    std::cout << "Failed to align pointclouds, exiting...";
                    break;
                }
                std::cout << "\nYou can find the scanned object parts in the output[Nr].ply files.";
            }

            break;
        case 3:
            pointclouds.clear();
            std::cout << "Scan buffer cleared.";
            wipe_old_scans();
            break;
        case 4:
            std::cout << "Exiting...";
            running = false;
            break;
        default:
            std::cout << "Invalid option, please try again.";
            break;
        }
    }

    return EXIT_SUCCESS;
}
catch (const rs2::error &e)
{
    std::cerr << "RealSense error calling " << e.get_failed_function() << "(" << e.get_failed_args() << "):\n    " << e.what() << std::endl;
    return EXIT_FAILURE;
}
catch (const std::exception &e)
{
    std::cerr << e.what() << std::endl;
    return EXIT_FAILURE;
}

void configure_filters(rs2::decimation_filter &dec_filter, rs2::threshold_filter &thr_filter, rs2::spatial_filter &spat_filter, int dec_filter_magnitude, float min_dist, float max_dist, int magnitude, int smooth_alpha, int smooth_delta, int holes_fill)
{
    dec_filter.set_option(RS2_OPTION_FILTER_MAGNITUDE, dec_filter_magnitude);
    thr_filter.set_option(RS2_OPTION_MIN_DISTANCE, min_dist);
    thr_filter.set_option(RS2_OPTION_MAX_DISTANCE, max_dist);
    spat_filter.set_option(RS2_OPTION_FILTER_MAGNITUDE, magnitude);
    spat_filter.set_option(RS2_OPTION_FILTER_SMOOTH_ALPHA, smooth_alpha);
    spat_filter.set_option(RS2_OPTION_FILTER_SMOOTH_DELTA, smooth_delta);
    spat_filter.set_option(RS2_OPTION_HOLES_FILL, holes_fill);
}

/* Apply filters.
    The implemented flow of the filters pipeline is in the following order:
    1. apply decimation filter
    2. apply threshold filter
    3. transform the scene into disparity domain
    4. apply spatial filter
    5. apply temporal filter
    6. revert the results back (if step Disparity filter was applied
    to depth domain (each post processing block is optional and can be applied independantly).
    */
void apply_filters(rs2::decimation_filter &dec_filter, rs2::threshold_filter &thr_filter, rs2::spatial_filter &spat_filter, rs2::disparity_transform &depth_to_disparity, rs2::disparity_transform &disparity_to_depth, rs2::depth_frame &depth_frame)
{
    depth_frame = dec_filter.process(depth_frame);
    depth_frame = thr_filter.process(depth_frame);
    depth_frame = depth_to_disparity.process(depth_frame);
    depth_frame = spat_filter.process(depth_frame);
    depth_frame = disparity_to_depth.process(depth_frame);
}

ColoredPoints scan(rs2::decimation_filter &dec_filter, rs2::threshold_filter &thr_filter, rs2::spatial_filter &spat_filter, rs2::disparity_transform &depth_to_disparity, rs2::disparity_transform &disparity_to_depth, rs2::pipeline &pipe)
{
    rs2::pointcloud pointcloud;
    rs2::points points;

    // let the auto exposure settle
    for (auto i = 0; i < 30; ++i)
        pipe.wait_for_frames();

    rs2::frameset data = pipe.wait_for_frames();
    rs2::video_frame color = data.get_color_frame();
    if (!color)
        color = data.get_infrared_frame();
    rs2::depth_frame depth_frame = data.get_depth_frame();

    apply_filters(dec_filter, thr_filter, spat_filter, depth_to_disparity, disparity_to_depth, depth_frame);

    pointcloud.map_to(color);
    points = pointcloud.calculate(depth_frame);

    return ColoredPoints(points, color);
}

void scan_and_convert(rs2::decimation_filter &dec_filter, rs2::threshold_filter &thr_filter, rs2::spatial_filter &spat_filter, rs2::disparity_transform &depth_to_disparity, rs2::disparity_transform &disparity_to_depth, rs2::pipeline &pipe, std::vector<PCLHelper::PtrCloud> &pointclouds)
{
    // does a scan, convertes it to pcl format and puts it into an array
    std::cout << "Scanning... ";
    ColoredPoints scanned_points = scan(dec_filter, thr_filter, spat_filter, depth_to_disparity, disparity_to_depth, pipe);
    std::cout << "\nConverting... ";
    PCLHelper::PtrCloud converted_points = PCLHelper::points_to_pcl(scanned_points.points, scanned_points.color);
    pointclouds.push_back(converted_points);
}

bool align_and_save_pointclouds(std::vector<PCLHelper::PtrCloud> &pointclouds)
{
    for (int i = 0; i < pointclouds.size(); i++)
    {
        if (i == 0)
        {
            PCLHelper::save_to_ply("scans_fitted/", "output" + std::to_string(i) + ".ply", pointclouds[i]);
            continue;
        }

        if (PCLHelper::icp(pointclouds[0], pointclouds[i], 80) == 0)
        {
            PCLHelper::save_to_ply("scans_fitted/", "output" + std::to_string(i) + ".ply", pointclouds[i]);
        }
        else
        {
            std::cout << "ICP failed";
            return false;
        }
    }
    return true;
}

void save_original_pointclouds(std::vector<PCLHelper::PtrCloud> &pointclouds)
{
    for (int i = 0; i < pointclouds.size(); i++)
    {
        PCLHelper::save_to_ply("scans_original/", "output" + std::to_string(i) + ".ply", pointclouds[i]);
    }
}
void wipe_old_scans()
{
    std::filesystem::path dirPathFitted = "scans_fitted";
    std::filesystem::path dirPathOriginal = "scans_original";

    std::uintmax_t numRemovedFitted = std::filesystem::remove_all(dirPathFitted);
    std::uintmax_t numRemovedOriginal = std::filesystem::remove_all(dirPathOriginal);

    if (numRemovedFitted == 0 && numRemovedOriginal == 0)
    {
        std::cout << "No Scans to delete. " << std::filesystem::current_path() << std::endl;
    }
    else
    {
        std::cout << "All scans delted." << std::endl;
    }
}

int prompt_user()
{
    int number;
    std::cout << "\n Welcome to the 3d-scanner, choose your option: ";
    std::cout << "\n (1) Scan ";
    std::cout << "\n (2) Align and save pointclouds ";
    std::cout << "\n (3) Wipe old scans ";
    std::cout << "\n (4) exit \n";
    std::cin >> number;
    return number;
}