/* +------------------------------------------------------------------------+
   |                     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          |
   +------------------------------------------------------------------------+ */

/* ------------------------------------------------------
          rgbd2rawlog
  A small tool to translate RGB+D datasets from TUM:
  http://cvpr.in.tum.de/data/datasets/rgbd-dataset/

  into MRPT rawlog binary format, ready to be parsed
  by RawLogViewer or user programs.

Usage:
    rgbd_dataset2rawlog  [PATH_TO_RGBD_DATASET] [OUTPUT_NAME]

Where expected input files are:
  - PATH_TO_RGBD_DATASET/accelerometer.txt
  - PATH_TO_RGBD_DATASET/depth.txt
  - PATH_TO_RGBD_DATASET/rgb.txt
  - PATH_TO_RGBD_DATASET/rgb/... (Images, 3x8u)
  - PATH_TO_RGBD_DATASET/depth/...  (Images, 1x16u)

Output files:
  - OUTPUT_NAME.rawlog: The output rawlog file.
  - OUTPUT_NAME_Images/...: External RGB images.

   ------------------------------------------------------ */

#include <mrpt/io/CFileGZOutputStream.h>
#include <mrpt/io/CTextFileLinesParser.h>
#include <mrpt/obs/CObservation3DRangeScan.h>
#include <mrpt/obs/CObservationIMU.h>
#include <mrpt/serialization/CArchive.h>
#include <mrpt/system/filesystem.h>

#include <iostream>

using namespace std;
using namespace mrpt;
using namespace mrpt::obs;
using namespace mrpt::serialization;

const double KINECT_FPS = 30.0;

// Based on:
// http://stackoverflow.com/questions/218488/finding-best-matching-key-for-a-given-key-in-a-sorted-stl-container
template <class CONTAINER>
typename CONTAINER::const_iterator find_closest(
    const CONTAINER& data, const typename CONTAINER::key_type& searchkey)
{
  typename CONTAINER::const_iterator upper = data.lower_bound(searchkey);
  if (upper == data.begin() || upper->first == searchkey) return upper;
  typename CONTAINER::const_iterator lower = upper;
  --lower;
  if (upper == data.end() || (searchkey - lower->first) < (upper->first - searchkey)) return lower;
  return upper;
}

void rgbd2rawlog(const string& src_path, const string& out_name)
{
  const string in_fil_acc = src_path + string("/accelerometer.txt");
  const string in_fil_depth = mrpt::system::fileExists(src_path + string("/ir.txt"))
                                  ? src_path + string("/ir.txt")
                                  : src_path + string("/depth.txt");
  const string in_fil_rgb = src_path + string("/rgb.txt");

  // ASSERT_FILE_EXISTS_(in_fil_acc)
  // ASSERT_FILE_EXISTS_(in_fil_depth)
  // ASSERT_FILE_EXISTS_(in_fil_rgb)

  // make a list with RGB & DEPTH files ------------------------
  map<double, string> list_rgb, list_depth;
  map<double, vector<double>> list_acc;
  std::istringstream line;
  if (mrpt::system::fileExists(in_fil_rgb))
  {
    mrpt::io::CTextFileLinesParser fparser(in_fil_rgb);
    while (fparser.getNextLine(line))
    {
      double tim;
      string f;
      if (line >> tim >> f) list_rgb[tim] = f;
    }
  }
  if (mrpt::system::fileExists(in_fil_depth))
  {
    mrpt::io::CTextFileLinesParser fparser(in_fil_depth);
    while (fparser.getNextLine(line))
    {
      double tim;
      string f;
      if (line >> tim >> f) list_depth[tim] = f;
    }
  }
  if (mrpt::system::fileExists(in_fil_acc))
  {
    mrpt::io::CTextFileLinesParser fparser(in_fil_acc);
    while (fparser.getNextLine(line))
    {
      double tim;
      double ax, ay, az;
      if (line >> tim >> ax >> ay >> az)
      {
        vector<double>& v = list_acc[tim];
        v.resize(3);
        v[0] = ax;
        v[1] = ay;
        v[2] = az;
      }
    }
  }
  cout << "Parsed: " << list_depth.size() << " / " << list_rgb.size() << " / " << list_acc.size()
       << " depth/rgb/acc entries.\n";

  const bool only_ir = list_depth.size() > 10 && list_rgb.empty();
  const bool only_rgb = list_depth.empty() && list_rgb.size() > 10;

  // Create output directory for images ------------------------------
  const string out_img_dir = out_name + string("_Images");

  cout << "Creating images directory: " << out_img_dir << endl;
  mrpt::system::createDirectory(out_img_dir);

  // Create rawlog file ----------------------------------------------
  const string out_rawlog_fil = out_name + string(".rawlog");
  cout << "Creating rawlog: " << out_rawlog_fil << endl;
  mrpt::io::CFileGZOutputStream f_out(out_rawlog_fil);

  // Fill out the common field to all entries:
  CObservation3DRangeScan obs;
  obs.sensorLabel = "KINECT";
  obs.range_is_depth = true;  // Kinect style: ranges are actually depth
  // values, not Euclidean distances.

  // Range & RGB images are already registered and warped to remove
  // distortion:
  const double FOCAL = 525.0;
  obs.cameraParams.nrows = 480;
  obs.cameraParams.ncols = 640;
  obs.cameraParams.fx(FOCAL);
  obs.cameraParams.fy(FOCAL);
  obs.cameraParams.cx(319.5);
  obs.cameraParams.cy(239.5);
  obs.cameraParamsIntensity = obs.cameraParams;
  obs.relativePoseIntensityWRTDepth = mrpt::poses::CPose3D(
      0, 0, 0, mrpt::DEG2RAD(-90), 0,
      mrpt::DEG2RAD(-90));  // No translation between rgb & range
  // cameras, and rotation only due to XYZ
  // axes conventions.

  CObservationIMU obs_imu;
  obs_imu.sensorLabel = "KINECT_ACC";
  obs_imu.dataIsPresent[IMU_X_ACC] = true;
  obs_imu.dataIsPresent[IMU_Y_ACC] = true;
  obs_imu.dataIsPresent[IMU_Z_ACC] = true;

  // Go thru the data:
  unsigned int counter = 0;
  if (!only_ir && !only_rgb)
  {
    for (map<double, string>::const_iterator it_list_rgb = list_rgb.begin();
         it_list_rgb != list_rgb.end(); ++it_list_rgb)
    {
      cout << "Processing image " << counter << "\r";
      cout.flush();
      counter++;

      // This is not the most efficient solution in the world, but...come
      // on! it's just a rawlog converter tool!
      map<double, string>::const_iterator it_list_depth =
          find_closest(list_depth, it_list_rgb->first);

      const double At = std::abs(it_list_rgb->first - it_list_depth->first);
      if (At > (1. / KINECT_FPS) * .5)
      {
        cout << "\nWarning: Discarding observation for too separated "
                "RGB/D timestamps: "
             << At * 1e3 << " ms\n";
      }
      else
      {
        // OK, we accept this RGB-DEPTH pair:
        const double avrg_time = .5 * (it_list_rgb->first + it_list_depth->first);
        obs.timestamp = mrpt::Clock::fromDouble(avrg_time);

        // RGB img:
        obs.hasIntensityImage = true;
        bool loadOk = obs.intensityImage.loadFromFile(src_path + string("/") + it_list_rgb->second);
        ASSERT_(loadOk);

        const string sRGBfile = mrpt::format("%.06f_rgb.png", avrg_time);
        bool savedOk = obs.intensityImage.saveToFile(out_img_dir + string("/") + sRGBfile);
        ASSERT_(savedOk);

        obs.intensityImage.setExternalStorage(sRGBfile);

        // Depth:
        obs.hasRangeImage = true;
        obs.rangeImage_forceResetExternalStorage();
        mrpt::img::CImage depth_img;
        if (!depth_img.loadFromFile(src_path + string("/") + it_list_depth->second))
          throw std::runtime_error(string("Error loading depth image!: ") + it_list_depth->second);

        const unsigned int w = depth_img.getWidth();
        const unsigned int h = depth_img.getHeight();
        obs.rangeImage_setSize(h, w);

        for (unsigned int row = 0; row < h; row++)
        {
          const uint16_t* ptr = depth_img.ptrLine<uint16_t>(row);
          for (unsigned int col = 0; col < w; col++)
            obs.rangeImage(row, col) = (*ptr++) * (1.f / 5000);
        }

        const string sDepthfile = mrpt::format("%.06f_depth.bin", avrg_time);
        obs.rangeImage_convertToExternalStorage(sDepthfile, out_img_dir + string("/"));

        // save:
        archiveFrom(f_out) << obs;

        // Search for acc data:
        map<double, vector<double>>::const_iterator it_list_acc = find_closest(list_acc, avrg_time);
        const double At_acc = std::abs(it_list_rgb->first - it_list_acc->first);
        if (At_acc < 1e-2)
        {
          obs_imu.timestamp = mrpt::Clock::fromDouble(avrg_time);

          obs_imu.rawMeasurements[IMU_X_ACC] = it_list_acc->second[0];
          obs_imu.rawMeasurements[IMU_Y_ACC] = it_list_acc->second[1];
          obs_imu.rawMeasurements[IMU_Z_ACC] = it_list_acc->second[2];

          archiveFrom(f_out) << obs_imu;
        }
      }
    }
  }
  else if (only_ir)
  {
    for (map<double, string>::const_iterator it_list_depth = list_depth.begin();
         it_list_depth != list_depth.end(); ++it_list_depth)
    {
      cout << "Processing image " << counter << "\r";
      cout.flush();
      counter++;

      const double avrg_time = it_list_depth->first;
      obs.timestamp = mrpt::Clock::fromDouble(it_list_depth->first);

      // Depth:
      obs.hasRangeImage = true;
      obs.rangeImage_forceResetExternalStorage();
      mrpt::img::CImage depth_img;
      if (!depth_img.loadFromFile(src_path + string("/") + it_list_depth->second))
        throw std::runtime_error(string("Error loading depth image!: ") + it_list_depth->second);

      const unsigned int w = depth_img.getWidth();
      const unsigned int h = depth_img.getHeight();
      obs.rangeImage_setSize(h, w);

      for (unsigned int row = 0; row < h; row++)
      {
        const uint16_t* ptr = depth_img.ptrLine<uint16_t>(row);
        for (unsigned int col = 0; col < w; col++)
          obs.rangeImage(row, col) = (*ptr++) * (1.f / 5000);
      }

      const string sDepthfile = mrpt::format("%.06f_depth.bin", avrg_time);
      obs.rangeImage_convertToExternalStorage(sDepthfile, out_img_dir + string("/"));

      // save:
      archiveFrom(f_out) << obs;
    }
  }
  else if (only_rgb)
  {
    for (map<double, string>::const_iterator it_list_rgb = list_rgb.begin();
         it_list_rgb != list_rgb.end(); ++it_list_rgb)
    {
      cout << "Processing image " << counter << "\r";
      cout.flush();
      counter++;

      const double avrg_time = it_list_rgb->first;
      obs.timestamp = mrpt::Clock::fromDouble(it_list_rgb->first);

      // RGB img:
      obs.hasIntensityImage = true;
      bool loadOk = obs.intensityImage.loadFromFile(src_path + string("/") + it_list_rgb->second);
      ASSERT_(loadOk);

      const string sRGBfile = mrpt::format("%.06f_rgb.png", avrg_time);
      bool savedOk = obs.intensityImage.saveToFile(out_img_dir + string("/") + sRGBfile);
      ASSERT_(savedOk);
      obs.intensityImage.setExternalStorage(sRGBfile);

      // save:
      archiveFrom(f_out) << obs;
    }
  }

  cout << "\nAll done!\n";
}

// ------------------------------------------------------
//						MAIN
// ------------------------------------------------------
int main(int argc, char** argv)
{
  try
  {
    if (argc != 3)
    {
      cerr << "Usage: " << argv[0] << " [PATH_TO_RGBD_DATASET] [OUTPUT_NAME]\n";
      return 1;
    }

    rgbd2rawlog(argv[1], argv[2]);

    return 0;
  }
  catch (const std::exception& e)
  {
    std::cerr << "MRPT error: " << mrpt::exception_to_str(e) << std::endl;
    return -1;
  }
}
