// ZED include
#include "ClientPublisher.hpp"
#include "GLViewer.hpp"
#include "utils.h"
#include "o3dutils.h"
#include "cuda_utils.h"

#include <open3d/Open3D.h>
#include <sl/Camera.hpp>

#include <mutex>
#include <thread>
#include <chrono>
#include <memory>
#include <iostream>
#include <cmath>

using namespace std;


int main(int argc, char **argv) 
{
	if (argc != 2) 
    {
		// this file should be generated by using the tool ZED360
		std::cout << "Need a Configuration file in input" << std::endl;		
		return 1;
	}
    
    // Defines the Coordinate system and unit used in this sample
    constexpr sl::COORDINATE_SYSTEM COORDINATE_SYSTEM = sl::COORDINATE_SYSTEM::RIGHT_HANDED_Y_UP;
    constexpr sl::UNIT UNIT = sl::UNIT::METER;

    // Read json file containing the configuration of your multicamera setup.    
    auto configurations = sl::readFusionConfigurationFile(argv[1], COORDINATE_SYSTEM, UNIT);

    if (configurations.empty()) 
    {
        std::cout << "Empty configuration File." << std::endl;
        return EXIT_FAILURE;
    }

    Trigger trigger;

    // Check if the ZED camera should run within the same process or if they are running on the edge.
    std::vector<ClientPublisher> clients(configurations.size());
    int id_ = 0;
    int gpu_id = 0;
    int nb_gpu = 0;
    cudaGetDeviceCount(&nb_gpu);

    for (auto conf: configurations) 
    {
        // if the ZED camera should run locally, then start a thread to handle it
        if(conf.communication_parameters.getType() == sl::CommunicationParameters::COMM_TYPE::INTRA_PROCESS)
        {
            std::cout << "Try to open ZED " <<conf.serial_number << ".." << std::flush;
            gpu_id = id_ % nb_gpu;
            auto state = clients[id_].open(conf.input_type, &trigger, gpu_id);
            if (!state) 
            {
                std::cerr << "Could not open ZED: " << conf.input_type.getConfiguration() << ". Skipping..." << std::endl;
                continue;
            }
            std::cout << ". ready !" << std::endl;
            id_++;
        }
    }

    // start camera threads - launch threads for each camera to begin data capture + streaming
    for (auto &it: clients)
        it.start();

    // Now that the ZED camera are running, we need to initialize the fusion module
    sl::InitFusionParameters init_params;
    init_params.coordinate_units = UNIT;
    init_params.coordinate_system = COORDINATE_SYSTEM;
    init_params.verbose = true;

    sl::Resolution low_res(1920, 1080);      // HD1080
    //sl::Resolution low_res(1280, 720);      // HD720
    init_params.maximum_working_resolution = low_res;
    // create and initialize it
    sl::Fusion fusion;
    fusion.init(init_params);

    // subscribe to every cameras of the setup to internally gather their data
    std::vector<sl::CameraIdentifier> cameras;
    for (auto& it : configurations) 
    {
        sl::CameraIdentifier uuid(it.serial_number);
        // to subscribe to a camera you must give its serial number, the way to communicate with it (shared memory or local network), and its world pose in the setup.        
        auto state = fusion.subscribe(uuid, it.communication_parameters, it.pose, it.override_gravity);
        if (state != sl::FUSION_ERROR_CODE::SUCCESS)
            std::cout << "Unable to subscribe to " << std::to_string(uuid.sn) << " . " << state << std::endl;
        else
            cameras.push_back(uuid);
    }

    // check that at least one camera is connected
    if (cameras.empty()) 
    {
        std::cout << "no connections " << std::endl;
        return EXIT_FAILURE;
    }

    // creation of a 3D viewer
    GLViewer viewer;
    viewer.init(argc, argv);
    cout << "Initialized OpenGL Viewer!\n";

    std::cout << "Viewer Shortcuts\n" <<
        "\t- 'q': quit the application\n" <<
        "\t- 'r': switch on/off for raw skeleton display\n" <<
        "\t- 'p': switch on/off for live point cloud display\n" <<
        "\t- 'c': switch on/off point cloud display with raw color\n" << std::endl;

    // fusion outputs
    std::map<sl::CameraIdentifier, sl::Bodies> camera_raw_data;
    std::map<sl::CameraIdentifier, sl::Mat> views;
    std::map<sl::CameraIdentifier, sl::Mat> pointClouds;
    std::mutex mutex_input;

    sl::CameraIdentifier fused_camera(0);
    sl::Mat fused_pc;
    sl::FusionMetrics metrics;

    while (viewer.isAvailable()) 
    {
        trigger.notifyZED();

        auto fusion_start = chrono::high_resolution_clock::now();
        auto fusion_code = fusion.process();
        auto fusion_end = chrono::high_resolution_clock::now();
        auto fusion_duration = chrono::duration<double, milli>(fusion_end - fusion_start).count();
        cout << "Total fusion time = " << fusion_duration << " ms." << endl;

        if (fusion_code == sl::FUSION_ERROR_CODE::SUCCESS) 
        {
            
            // for debug, you can retrieve the data sent by each camera
            for (auto& id : cameras) 
            { 
                sl::Pose pose;
                if(fusion.getPosition(pose, sl::REFERENCE_FRAME::WORLD, id, sl::POSITION_TYPE::RAW) == sl::POSITIONAL_TRACKING_STATE::OK)
                    viewer.setCameraPose(id.sn, pose.pose_data);
                    //viewer.setCameraPose(fused_camera.sn, pose.pose_data);

                auto state_view = fusion.retrieveImage(views[id], id, low_res);
                auto state_pc = fusion.retrieveMeasure(pointClouds[id], id, sl::MEASURE::XYZBGRA, low_res);

                //auto state_view = fusion.retrieveImage(views[fused_camera], fused_camera, low_res);
                //auto state_pc = fusion.retrieveMeasure(pointClouds[fused_camera], fused_camera, sl::MEASURE::XYZBGRA, low_res);

                if (state_view == sl::FUSION_ERROR_CODE::SUCCESS && state_pc == sl::FUSION_ERROR_CODE::SUCCESS)
                    viewer.updateFusion(id.sn, views[id], pointClouds[id]);
            }
            

            mutex_input.lock();
            auto fused_pc_status = fusion.retrieveMeasure(fused_pc, fused_camera, sl::MEASURE::XYZBGRA, low_res, sl::FUSION_REFERENCE_FRAME::WORLD);
            
            if (fused_pc_status == sl::FUSION_ERROR_CODE::SUCCESS) 
            {
                cout << "fused_pc status = " << sl::toString(fused_pc_status) << endl;

                auto pc_start = chrono::high_resolution_clock::now();

                auto o3d_pc = slMat_to_Open3D(fused_pc);

                // Voxel Downsample
                double voxel_size = 0.01; 
                auto voxelized_pcd = o3d_pc.VoxelDownSample(voxel_size);
                auto updated = Open3D_to_slMat(*voxelized_pcd);

                // viewer.updatePC(fused_camera.sn, updated);
                
                auto pc_end = chrono::high_resolution_clock::now();
                auto pc_duration = chrono::duration_cast<chrono::milliseconds>(pc_end - pc_start).count();
                cout << "Point cloud processing time (conversion + saving + voxel filtering) = " << pc_duration << " ms." << endl;

                std::string raw_filename = "../fused_raw_point_cloud.ply";
                std::string voxel_filename = "../fused_voxel_point_cloud.ply";

                // Saving the fused point cloud
                open3d::io::WritePointCloud(raw_filename, o3d_pc);
                cout << "Saved fused point cloud with " << o3d_pc.points_.size() << " points." << endl;

                // Save voxelized fused point cloud
                open3d::io::WritePointCloud(voxel_filename, *voxelized_pcd);
                cout << "Saved voxel filtered point cloud with " << voxelized_pcd->points_.size() << " points." << endl;
                
                cout << "------------------------------------------------------------------" << endl;
            
            }
            // get metrics about the fusion process for monitoring purposes
            fusion.getProcessMetrics(metrics);
            mutex_input.unlock();
        }
        std::this_thread::sleep_for(std::chrono::milliseconds(10));
    }

    trigger.running = false;
    trigger.notifyZED();

    // Stop viewer before clients
    viewer.exit(); 

    for (auto &it: clients)
        it.stop();

    fusion.close();

    return EXIT_SUCCESS;
}