/* +------------------------------------------------------------------------+
   |                     Mobile Robot Programming Toolkit (MRPT)            |
   |                          https://www.mrpt.org/                         |
   |                                                                        |
   | Copyright (c) 2005-2024, Individual contributors, see AUTHORS file     |
   | See: https://www.mrpt.org/Authors - All rights reserved.               |
   | Released under BSD License. See: https://www.mrpt.org/License          |
   +------------------------------------------------------------------------+ */

/*
  Example  : kinect-to-2d-laser-demo
  Web page : https://www.mrpt.org/Example_Kinect_To_2D_laser_scan (includes
  video demo)

  Purpose  : Demonstrate grabbing from CKinect, multi-threading
       and converting the 3D range data into an equivalent
       2D planar scan.
*/

#include <mrpt/config/CConfigFile.h>
#include <mrpt/gui/CDisplayWindow3D.h>
#include <mrpt/hwdrivers/CKinect.h>
#include <mrpt/img/TColor.h>
#include <mrpt/opengl/CFrustum.h>
#include <mrpt/opengl/CGridPlaneXY.h>
#include <mrpt/opengl/CPlanarLaserScan.h>
#include <mrpt/opengl/CPointCloudColoured.h>
#include <mrpt/opengl/stock_objects.h>
#include <mrpt/system/CTicTac.h>
#include <mrpt/system/filesystem.h>

#include <chrono>
#include <iostream>
#include <thread>

using namespace mrpt;
using namespace mrpt::hwdrivers;
using namespace mrpt::math;
using namespace mrpt::gui;
using namespace mrpt::obs;
using namespace mrpt::maps;
using namespace mrpt::opengl;
using namespace mrpt::system;
using namespace mrpt::img;
using namespace std;

// Thread for grabbing: Do this is another thread so we divide rendering and
// grabbing
//   and exploit multicore CPUs.
struct TThreadParam
{
  TThreadParam() {}
  volatile bool quit{false};
  volatile int pushed_key{0};
  volatile double Hz{0};

  CObservation3DRangeScan::Ptr new_obs;  // RGB+D (+3D points)
  CObservationIMU::Ptr new_obs_imu;      // Accelerometers
};

void thread_grabbing(TThreadParam& p)
{
  try
  {
    CKinect kinect;

    // Set params:
    // kinect.enableGrab3DPoints(true);
    // kinect.enablePreviewRGB(true);
    //...
    const std::string cfgFile = "kinect_calib.cfg";
    if (mrpt::system::fileExists(cfgFile))
    {
      cout << "Loading calibration from: " << cfgFile << endl;
      kinect.loadConfig(mrpt::config::CConfigFile(cfgFile), "KINECT");
    }
    else
      cerr << "Warning: Calibration file [" << cfgFile << "] not found -> Using default params.\n";

    // Open:
    cout << "Calling CKinect::initialize()...";
    kinect.initialize();
    cout << "OK\n";

    CTicTac tictac;
    int nImgs = 0;
    bool there_is_obs = true, hard_error = false;

    while (!hard_error && !p.quit)
    {
      // Grab new observation from the camera:
      auto obs = CObservation3DRangeScan::Create();  // Smart pointers to
                                                     // observations
      CObservationIMU::Ptr obs_imu = CObservationIMU::Create();

      kinect.getNextObservation(*obs, *obs_imu, there_is_obs, hard_error);

      if (!hard_error && there_is_obs)
      {
        std::atomic_store(&p.new_obs, obs);
        std::atomic_store(&p.new_obs_imu, obs_imu);
      }

      if (p.pushed_key != 0)
      {
        switch (p.pushed_key)
        {
          case 27:
            p.quit = true;
            break;
        }

        // Clear pushed key flag:
        p.pushed_key = 0;
      }

      nImgs++;
      if (nImgs > 10)
      {
        p.Hz = nImgs / tictac.Tac();
        nImgs = 0;
        tictac.Tic();
      }
    }
  }
  catch (const std::exception& e)
  {
    cout << "Exception in Kinect thread: " << e.what() << endl;
    p.quit = true;
  }
}

// ------------------------------------------------------
//				Test_Kinect
// ------------------------------------------------------
void Test_Kinect()
{
  // Launch grabbing thread:
  // --------------------------------------------------------
  TThreadParam thrPar;
  std::thread thHandle(thread_grabbing, std::ref(thrPar));

  // Wait until data stream starts so we can say for sure the sensor has been
  // initialized OK:
  cout << "Waiting for sensor initialization...\n";
  do
  {
    CObservation3DRangeScan::Ptr possiblyNewObs = std::atomic_load(&thrPar.new_obs);
    if (possiblyNewObs && possiblyNewObs->timestamp != INVALID_TIMESTAMP)
      break;
    else
      std::this_thread::sleep_for(10ms);
  } while (!thrPar.quit);

  // Check error condition:
  if (thrPar.quit) return;

  // Create window and prepare OpenGL object in the scene:
  // --------------------------------------------------------
  mrpt::gui::CDisplayWindow3D win3D("Kinect 3D -> 2D laser scan", 800, 600);

  win3D.setCameraAzimuthDeg(140);
  win3D.setCameraElevationDeg(30);
  win3D.setCameraZoom(10.0);
  win3D.setFOV(90);
  win3D.setCameraPointingToPoint(2.5, 0, 0);

  // The 3D point cloud OpenGL object:
  mrpt::opengl::CPointCloudColoured::Ptr gl_points = mrpt::opengl::CPointCloudColoured::Create();
  gl_points->setPointSize(2.5);

  // The 2D "laser scan" OpenGL object:
  mrpt::opengl::CPlanarLaserScan::Ptr gl_2d_scan = mrpt::opengl::CPlanarLaserScan::Create();
  gl_2d_scan->enablePoints(true);
  gl_2d_scan->enableLine(true);
  gl_2d_scan->enableSurface(true);
  gl_2d_scan->setSurfaceColor(0, 0, 1, 0.3);  // RGBA

  mrpt::opengl::CFrustum::Ptr gl_frustum =
      mrpt::opengl::CFrustum::Create(0.2f, 5.0f, 90.0f, 5.0f, 2.0f, true, true);

  const double aspect_ratio = 480.0 / 640.0;  // kinect.rows() / double( kinect.cols() );

  opengl::Viewport::Ptr viewRange,
      viewInt;  // Extra viewports for the RGB & D images.
  {
    mrpt::opengl::Scene::Ptr& scene = win3D.get3DSceneAndLock();

    // Create the Opengl object for the point cloud:
    scene->insert(gl_points);
    scene->insert(gl_2d_scan);
    scene->insert(gl_frustum);

    {
      mrpt::opengl::CGridPlaneXY::Ptr gl_grid = mrpt::opengl::CGridPlaneXY::Create();
      gl_grid->setColor(0.6, 0.6, 0.6);
      scene->insert(gl_grid);
    }
    {
      mrpt::opengl::CSetOfObjects::Ptr gl_corner = mrpt::opengl::stock_objects::CornerXYZ();
      gl_corner->setScale(0.2);
      scene->insert(gl_corner);
    }

    const int VW_WIDTH = 250;  // Size of the viewport into the window, in pixel units.
    const int VW_HEIGHT = aspect_ratio * VW_WIDTH;
    const int VW_GAP = 30;

    // Create the Opengl objects for the planar images, as textured planes,
    // each in a separate viewport:
    win3D.addTextMessage(30, -25 - 1 * (VW_GAP + VW_HEIGHT), "Range data", 1);
    viewRange = scene->createViewport("view2d_range");
    viewRange->setViewportPosition(5, -10 - 1 * (VW_GAP + VW_HEIGHT), VW_WIDTH, VW_HEIGHT);

    win3D.addTextMessage(30, -25 - 2 * (VW_GAP + VW_HEIGHT), "Intensity data", 2);
    viewInt = scene->createViewport("view2d_int");
    viewInt->setViewportPosition(5, -10 - 2 * (VW_GAP + VW_HEIGHT), VW_WIDTH, VW_HEIGHT);

    win3D.unlockAccess3DScene();
    win3D.repaint();
  }

  CObservation3DRangeScan::Ptr last_obs;
  CObservationIMU::Ptr last_obs_imu;

  while (win3D.isOpen() && !thrPar.quit)
  {
    CObservation3DRangeScan::Ptr possiblyNewObs = std::atomic_load(&thrPar.new_obs);
    if (possiblyNewObs && possiblyNewObs->timestamp != INVALID_TIMESTAMP &&
        (!last_obs || possiblyNewObs->timestamp != last_obs->timestamp))
    {
      // It IS a new observation:
      last_obs = possiblyNewObs;
      last_obs_imu = std::atomic_load(&thrPar.new_obs_imu);

      // Update visualization ---------------------------------------
      bool do_refresh = false;

      // Show 2D ranges as a grayscale image:
      if (last_obs->hasRangeImage)
      {
        mrpt::img::CImage img;

        // Normalize the image
        static CMatrixFloat range2D;  // Static to save time allocating
        // the matrix in every iteration
        range2D = last_obs->rangeImage;
        range2D *= (1.0 / last_obs->maxRange);

        img.setFromMatrix(range2D);

        win3D.get3DSceneAndLock();
        viewRange->setImageView(std::move(img));
        win3D.unlockAccess3DScene();
        do_refresh = true;
      }

      // Convert ranges to an equivalent 2D "fake laser" scan:
      if (last_obs->hasRangeImage)
      {
        // Convert to scan:
        CObservation2DRangeScan::Ptr obs_2d = CObservation2DRangeScan::Create();
        const float vert_FOV = DEG2RAD(gl_frustum->getVertFOV());

        mrpt::obs::T3DPointsTo2DScanParams sp;
        sp.angle_inf = .5f * vert_FOV;
        sp.angle_sup = .5f * vert_FOV;
        sp.sensorLabel = "KINECT_2D_SCAN";

        last_obs->convertTo2DScan(*obs_2d, sp);

        // And load scan in the OpenGL object:
        gl_2d_scan->setScan(*obs_2d);
      }

      // Show intensity image:
      if (last_obs->hasIntensityImage)
      {
        win3D.get3DSceneAndLock();
        viewInt->setImageView(last_obs->intensityImage);  // This is not "_fast" since the
        // intensity image is used below
        // in the coloured point cloud.
        win3D.unlockAccess3DScene();
        do_refresh = true;
      }

      // Show 3D points:
      if (last_obs->hasPoints3D)
      {
        win3D.get3DSceneAndLock();
        mrpt::obs::T3DPointsProjectionParams pp;
        pp.takeIntoAccountSensorPoseOnRobot = false;

        last_obs->unprojectInto(*gl_points, pp);
        win3D.unlockAccess3DScene();
        do_refresh = true;
      }

      // Some text messages:
      win3D.get3DSceneAndLock();
      // Estimated grabbing rate:
      win3D.addTextMessage(-100, -20, format("%.02f Hz", thrPar.Hz), 100);

      win3D.addTextMessage(
          10, 10,
          "'o'/'i'-zoom out/in, '2'/'3'/'f':show/hide 2D/3D/frustum "
          "data, mouse: orbit 3D, ESC: quit",
          110);

      win3D.addTextMessage(
          10, 25,
          format(
              "Show: 3D=%s 2D=%s Frustum=%s (vert. FOV=%.1fdeg)",
              gl_points->isVisible() ? "YES" : "NO", gl_2d_scan->isVisible() ? "YES" : "NO",
              gl_frustum->isVisible() ? "YES" : "NO", gl_frustum->getVertFOV()),
          111);
      win3D.unlockAccess3DScene();

      // Do we have accelerometer data?
      if (last_obs_imu && last_obs_imu->dataIsPresent[IMU_X_ACC])
      {
        win3D.get3DSceneAndLock();
        win3D.addTextMessage(
            10, 65,
            format(
                "Acc: x=%.02f y=%.02f z=%.02f", last_obs_imu->rawMeasurements[IMU_X_ACC],
                last_obs_imu->rawMeasurements[IMU_Y_ACC], last_obs_imu->rawMeasurements[IMU_Z_ACC]),
            102);
        win3D.unlockAccess3DScene();
        do_refresh = true;
      }

      // Force opengl repaint:
      if (do_refresh) win3D.repaint();

    }  // end update visualization:

    // Process possible keyboard commands:
    // --------------------------------------
    if (win3D.keyHit() && thrPar.pushed_key == 0)
    {
      const int key = tolower(win3D.getPushedKey());

      switch (key)
      {
        // Some of the keys are processed in this thread:
        case 'o':
          win3D.setCameraZoom(win3D.getCameraZoom() * 1.2);
          win3D.repaint();
          break;
        case 'i':
          win3D.setCameraZoom(win3D.getCameraZoom() / 1.2);
          win3D.repaint();
          break;
        case '2':
          gl_2d_scan->setVisibility(!gl_2d_scan->isVisible());
          break;
        case '3':
          gl_points->setVisibility(!gl_points->isVisible());
          break;
        case 'F':
        case 'f':
          gl_frustum->setVisibility(!gl_frustum->isVisible());
          break;
        case '+':
          gl_frustum->setVertFOV(gl_frustum->getVertFOV() + 1);
          break;
        case '-':
          gl_frustum->setVertFOV(gl_frustum->getVertFOV() - 1);
          break;
        // ...and the rest in the kinect thread:
        default:
          thrPar.pushed_key = key;
          break;
      };
    }

    std::this_thread::sleep_for(1ms);
  }

  cout << "Waiting for grabbing thread to exit...\n";
  thrPar.quit = true;
  thHandle.join();
  cout << "Bye!\n";
}

int main(int argc, char** argv)
{
  try
  {
    Test_Kinect();

    std::this_thread::sleep_for(50ms);
    return 0;
  }
  catch (const std::exception& e)
  {
    std::cerr << "MRPT error: " << mrpt::exception_to_str(e) << std::endl;
    return -1;
  }
  catch (...)
  {
    printf("Another exception!!");
    return -1;
  }
}
