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

#include "apps-precomp.h"  // Precompiled headers
//
#include <mrpt/3rdparty/tclap/CmdLine.h>
#include <mrpt/apps/RawlogGrabberApp.h>
#include <mrpt/config/CConfigFile.h>
#include <mrpt/core/lock_helper.h>
#include <mrpt/core/round.h>
#include <mrpt/hwdrivers/CGenericSensor.h>
#include <mrpt/img/CImage.h>
#include <mrpt/io/CFileGZOutputStream.h>
#include <mrpt/io/vector_loadsave.h>
#include <mrpt/obs/CActionCollection.h>
#include <mrpt/obs/CActionRobotMovement2D.h>
#include <mrpt/obs/CObservationGPS.h>
#include <mrpt/obs/CObservationIMU.h>
#include <mrpt/obs/CObservationOdometry.h>
#include <mrpt/obs/CSensoryFrame.h>
#include <mrpt/serialization/CArchive.h>
#include <mrpt/system/CRateTimer.h>
#include <mrpt/system/filesystem.h>
#include <mrpt/system/os.h>
#include <mrpt/system/thread_name.h>

#include <thread>

using namespace mrpt::apps;

RawlogGrabberApp::RawlogGrabberApp() : mrpt::system::COutputLogger("RawlogGrabberApp") {}

void RawlogGrabberApp::initialize(int argc, const char** argv)
{
  MRPT_START

  if ((getenv("MRPT_HWDRIVERS_VERBOSE") != nullptr) && atoi(getenv("MRPT_HWDRIVERS_VERBOSE")) != 0)
  {
    this->setMinLoggingLevel(mrpt::system::LVL_DEBUG);
  }

  MRPT_LOG_INFO(" rawlog-grabber - Part of the MRPT");
  MRPT_LOG_INFO_FMT(
      " MRPT C++ Library: %s - Sources timestamp: %s\n", mrpt::system::MRPT_getVersion().c_str(),
      mrpt::system::MRPT_getCompilationDate().c_str());

  // Declare the supported options.
  TCLAP::CmdLine cmd("rawlog-grabber", ' ', mrpt::system::MRPT_getVersion().c_str());

  TCLAP::UnlabeledValueArg<std::string> argConfigFile(
      "config", "Config file", true, "", "<configFile.ini>", cmd);
  TCLAP::ValueArg<std::string> argPlugins(
      "p", "plugins",
      "Load one or more plug-in modules (.so/.dll) with additional sensor "
      "drivers (comma-separated list)",
      false, "", "myModule.so", cmd);

  // Process arguments:
  if (!cmd.parse(argc, argv)) THROW_EXCEPTION("CLI arguments parsing tells we should exit.");

  ASSERT_FILE_EXISTS_(argConfigFile.getValue());
  params.setContent(mrpt::io::file_get_contents(argConfigFile.getValue()));

  if (argPlugins.isSet()) mrpt::system::loadPluginModules(argPlugins.getValue());

  MRPT_END
}

void RawlogGrabberApp::runImpl()
{
  using namespace mrpt;
  using namespace mrpt::system;
  using namespace mrpt::hwdrivers;
  using namespace mrpt::config;
  using namespace mrpt::serialization;
  using namespace mrpt::img;
  using namespace mrpt::obs;
  using namespace mrpt::poses;
  using namespace std;

  const std::string GLOBAL_SECT = "global";

  // ------------------------------------------
  //			Load config from file:
  // ------------------------------------------
  string rawlog_prefix = "dataset";
  int time_between_launches = 300;
  bool use_sensoryframes = false;
  int GRABBER_PERIOD_MS = 1000;
  int rawlog_GZ_compress_level = 1;  // 0: No compress, 1-9: compress level

  MRPT_LOAD_CONFIG_VAR(rawlog_prefix, string, params, GLOBAL_SECT);
  MRPT_LOAD_CONFIG_VAR(time_between_launches, int, params, GLOBAL_SECT);
  MRPT_LOAD_CONFIG_VAR(SF_max_time_span, float, params, GLOBAL_SECT);
  MRPT_LOAD_CONFIG_VAR(use_sensoryframes, bool, params, GLOBAL_SECT);
  MRPT_LOAD_CONFIG_VAR(GRABBER_PERIOD_MS, int, params, GLOBAL_SECT);

  MRPT_LOAD_CONFIG_VAR(rawlog_GZ_compress_level, int, params, GLOBAL_SECT);

  // Build full rawlog file name:
  string rawlog_postfix = "_";

  // rawlog_postfix += dateTimeToString( now() );
  mrpt::system::TTimeParts parts;
  mrpt::system::timestampToParts(mrpt::Clock::now(), parts, true);
  rawlog_postfix += format(
      "%04u-%02u-%02u_%02uh%02um%02us", (unsigned int)parts.year, (unsigned int)parts.month,
      (unsigned int)parts.day, (unsigned int)parts.hour, (unsigned int)parts.minute,
      (unsigned int)parts.second);

  rawlog_postfix = mrpt::system::fileNameStripInvalidChars(rawlog_postfix);

  // Only set this if we want externally stored images:
  m_rawlog_ext_imgs_dir =
      rawlog_prefix + fileNameStripInvalidChars(rawlog_postfix + string("_Images"));

  // Also, set the path in CImage to enable online visualization in a GUI
  // window:
  CImage::setImagesPathBase(m_rawlog_ext_imgs_dir);

  rawlog_postfix += string(".rawlog");
  rawlog_postfix = fileNameStripInvalidChars(rawlog_postfix);

  results_mtx.lock();
  rawlog_filename = rawlog_prefix + rawlog_postfix;
  results_mtx.unlock();

  MRPT_LOG_INFO_STREAM("Output rawlog filename: " << rawlog_filename);
  MRPT_LOG_INFO_STREAM("External image storage: " << m_rawlog_ext_imgs_dir);

  // ----------------------------------------------
  // Launch threads:
  // ----------------------------------------------
  std::vector<std::string> sections;
  params.getAllSections(sections);

  vector<std::thread> lstThreads;

  for (auto& section : sections)
  {
    if (section == GLOBAL_SECT || section.empty() ||
        params.read_bool(section, "rawlog-grabber-ignore", false, false))
      continue;  // This is not a sensor:

    std::thread& newThread =
        lstThreads.emplace_back(&RawlogGrabberApp::SensorThread, this, section);

    mrpt::system::thread_name(section, newThread);

    std::this_thread::sleep_for(std::chrono::milliseconds(time_between_launches));
  }

  // ----------------------------------------------
  // Run:
  // ----------------------------------------------
  mrpt::io::CFileGZOutputStream out_file;
  auto out_arch_obj = archiveFrom(out_file);
  m_out_arch_ptr = &out_arch_obj;

  out_file.open(rawlog_filename, rawlog_GZ_compress_level);

  CGenericSensor::TListObservations copy_of_m_global_list_obs;

  MRPT_LOG_INFO_STREAM("Press any key to exit program");

  mrpt::system::CTicTac run_timer;
  run_timer.Tic();

  auto lambdaProcessPending = [&]()
  {
    auto lock = mrpt::lockHelper(cs_m_global_list_obs);
    copy_of_m_global_list_obs.clear();

    if (!m_global_list_obs.empty())
    {
      auto itEnd = m_global_list_obs.begin();
      std::advance(itEnd, m_global_list_obs.size() / 2);
      copy_of_m_global_list_obs.insert(m_global_list_obs.begin(), itEnd);
      m_global_list_obs.erase(m_global_list_obs.begin(), itEnd);
    }

    if (use_sensoryframes)
      process_observations_for_sf(copy_of_m_global_list_obs);
    else
      process_observations_for_nonsf(copy_of_m_global_list_obs);
  };

  while (!os::kbhit() && !allThreadsMustExit())
  {
    // Check "run for X seconds" flag:
    {
      auto lk = mrpt::lockHelper(params_mtx);
      if (run_for_seconds > 0 && run_timer.Tac() > run_for_seconds) break;
    }

    // See if we have observations and process them:
    lambdaProcessPending();

    std::this_thread::sleep_for(std::chrono::milliseconds(GRABBER_PERIOD_MS));
  }

  if (allThreadsMustExit())
  {
    MRPT_LOG_ERROR(
        "[main thread] Ended due to other thread signal to exit "
        "application.");
  }

  // Final check of pending objects:
  lambdaProcessPending();

  // Flush file to disk:
  out_file.close();

  // Wait all threads:
  // ----------------------------
  allThreadsMustExit(true);
  std::this_thread::sleep_for(100ms);
  MRPT_LOG_INFO("Waiting for all threads to close...");
  for (auto& t : lstThreads)
    if (t.joinable()) t.join();
}

void RawlogGrabberApp::run()
{
  try
  {
    m_isRunningMtx.lock();
    m_isRunning = true;
    m_isRunningMtx.unlock();

    runImpl();

    m_isRunningMtx.lock();
    m_isRunning = false;
    m_isRunningMtx.unlock();
  }
  catch (const std::exception& e)
  {
    m_isRunningMtx.lock();
    m_isRunning = false;
    m_isRunningMtx.unlock();
    throw;
  }
}

void RawlogGrabberApp::dump_verbose_info(const mrpt::serialization::CSerializable::Ptr& o) const
{
  // Show GPS mode:
  if (getMinLoggingLevel() != mrpt::system::LVL_DEBUG) return;

  static auto last_t = mrpt::Clock::now();
  const auto t_now = mrpt::Clock::now();
  if (mrpt::system::timeDifference(last_t, t_now) < 1.0) return;
  last_t = t_now;

  if (auto gps = std::dynamic_pointer_cast<mrpt::obs::CObservationGPS>(o); gps)
  {
    dump_GPS_mode_info(*gps);
  }
  else if (auto imu = std::dynamic_pointer_cast<mrpt::obs::CObservationIMU>(o); imu)
  {
    dump_IMU_info(*imu);
  }
}

void RawlogGrabberApp::dump_verbose_info(const mrpt::obs::CSensoryFrame& sf) const
{
  if (getMinLoggingLevel() != mrpt::system::LVL_DEBUG) return;

  // Show GPS mode:
  mrpt::obs::CObservationGPS::Ptr gps;
  std::size_t idx = 0;
  do
  {
    gps = sf.getObservationByClass<mrpt::obs::CObservationGPS>(idx++);
    if (gps) dump_GPS_mode_info(*gps);
  } while (gps);

  // Show IMU angles:
  auto imu = sf.getObservationByClass<mrpt::obs::CObservationIMU>();
  if (imu) dump_IMU_info(*imu);
}

void RawlogGrabberApp::dump_GPS_mode_info(const mrpt::obs::CObservationGPS& o) const
{
  if (o.has_GGA_datum())
  {
    using mrpt::obs::gnss::Message_NMEA_GGA;

    const auto fq = static_cast<int>(o.getMsgByClass<Message_NMEA_GGA>().fields.fix_quality);
    MRPT_LOG_DEBUG_STREAM("  GPS mode: " << fq << " label: " << o.sensorLabel);
  }
  {
    std::stringstream ss;
    o.getDescriptionAsText(ss);
    MRPT_LOG_DEBUG_STREAM(ss.str());
  }
}

void RawlogGrabberApp::dump_IMU_info(const mrpt::obs::CObservationIMU& o) const
{
  // Show IMU angles:
  MRPT_LOG_DEBUG_FMT(
      "   IMU angles (degrees): "
      "(yaw,pitch,roll)=(%.06f, %.06f, %.06f)",
      mrpt::RAD2DEG(o.rawMeasurements[mrpt::obs::IMU_YAW]),
      mrpt::RAD2DEG(o.rawMeasurements[mrpt::obs::IMU_PITCH]),
      mrpt::RAD2DEG(o.rawMeasurements[mrpt::obs::IMU_ROLL]));
}

// ------------------------------------------------------
//				SensorThread
// ------------------------------------------------------
void RawlogGrabberApp::SensorThread(std::string sensor_label)
{
  try
  {
    std::string driver_name = params.read_string(sensor_label, "driver", "", true);

    auto sensor = mrpt::hwdrivers::CGenericSensor::createSensorPtr(driver_name);

    if (!sensor) throw std::runtime_error(std::string("Class name not recognized: ") + driver_name);

    // Load common & sensor specific parameters:
    sensor->loadConfig(params, sensor_label);

    MRPT_LOG_INFO_STREAM(
        "[thread_" << sensor_label << "] Starting at " << sensor->getProcessRate() << " Hz");

    ASSERT_GT_(sensor->getProcessRate(), 0);

    // For imaging sensors, set external storage directory:
    sensor->setPathForExternalImages(m_rawlog_ext_imgs_dir);

    // Init device:
    sensor->initialize();

    mrpt::system::CRateTimer rate;
    rate.setRate(sensor->getProcessRate());

    while (!allThreadsMustExit())
    {
      // Process
      sensor->doProcess();

      // Get new observations
      const auto lstObjs = sensor->getObservations();

      {
        std::lock_guard<std::mutex> lock(cs_m_global_list_obs);
        m_global_list_obs.insert(lstObjs.begin(), lstObjs.end());
      }

      // wait for the process period:
      rate.sleep();
    }

    sensor.reset();

    MRPT_LOG_INFO_FMT("[thread_%s] Closing...", sensor_label.c_str());
  }
  catch (const std::exception& e)
  {
    if (show_sensor_thread_exceptions)
    {
      MRPT_LOG_ERROR_STREAM("Exception in SensorThread:\n" << mrpt::exception_to_str(e));
    }
    allThreadsMustExit(true);
  }
  catch (...)
  {
    if (show_sensor_thread_exceptions)
    {
      MRPT_LOG_ERROR("Untyped exception in SensorThread.");
    }
    allThreadsMustExit(true);
  }
}

void RawlogGrabberApp::process_observations_for_sf(
    const RawlogGrabberApp::TListObservations& list_obs)
{
  using namespace mrpt::obs;

  // -----------------------
  // USE SENSORY-FRAMES
  // -----------------------
  for (auto it = list_obs.begin(); it != list_obs.end(); ++it)
  {
    // If we have an action, save the SF and start a new one:
    if (IS_DERIVED(*it->second, CAction))
    {
      CAction::Ptr act = std::dynamic_pointer_cast<CAction>(it->second);

      (*m_out_arch_ptr) << m_curSF;
      MRPT_LOG_INFO_STREAM("Saved SF with " << m_curSF.size() << " objects.");
      m_curSF.clear();

      CActionCollection acts;
      acts.insert(*act);
      act.reset();

      (*m_out_arch_ptr) << acts;
    }
    else if (IS_CLASS(*it->second, CObservationOdometry))
    {
      CObservationOdometry::Ptr odom = std::dynamic_pointer_cast<CObservationOdometry>(it->second);

      auto act = CActionRobotMovement2D::Create();
      act->timestamp = odom->timestamp;

      // Compute the increment since the last reading:
      static CActionRobotMovement2D::TMotionModelOptions odomOpts;
      static CObservationOdometry last_odo;
      static bool last_odo_first = true;

      mrpt::poses::CPose2D odo_incr;
      int64_t lticks_incr, rticks_incr;

      if (last_odo_first)
      {
        last_odo_first = false;
        odo_incr = mrpt::poses::CPose2D(0, 0, 0);
        lticks_incr = rticks_incr = 0;
      }
      else
      {
        odo_incr = odom->odometry - last_odo.odometry;
        lticks_incr = odom->encoderLeftTicks - last_odo.encoderLeftTicks;
        rticks_incr = odom->encoderRightTicks - last_odo.encoderRightTicks;

        last_odo = *odom;
      }

      // Save as action & dump to file:
      act->computeFromOdometry(odo_incr, odomOpts);

      act->hasEncodersInfo = true;
      act->encoderLeftTicks = lticks_incr;
      act->encoderRightTicks = rticks_incr;

      act->hasVelocities = true;
      act->velocityLocal = odom->velocityLocal;

      (*m_out_arch_ptr) << m_curSF;

      MRPT_LOG_INFO_STREAM("Saved SF with " << m_curSF.size() << " objects.");
      m_curSF.clear();

      CActionCollection acts;
      acts.insert(*act);
      act.reset();

      (*m_out_arch_ptr) << acts;
      {
        auto lk = mrpt::lockHelper(results_mtx);
        rawlog_saved_objects += 2;  // m_curSF + acts;
      }
    }
    else if (IS_DERIVED(*it->second, CObservation))
    {
      CObservation::Ptr obs = std::dynamic_pointer_cast<CObservation>(it->second);

      // First, check if inserting this OBS into the SF would
      // overflow "SF_max_time_span":
      if (m_curSF.size() != 0 &&
          mrpt::system::timeDifference(
              m_curSF.getObservationByIndex(0)->timestamp, obs->timestamp) > SF_max_time_span)
      {
        dump_verbose_info(m_curSF);

        // Save and start a new one:
        (*m_out_arch_ptr) << m_curSF;
        {
          auto lk = mrpt::lockHelper(results_mtx);
          rawlog_saved_objects++;
        }

        MRPT_LOG_INFO_STREAM("Saved SF with " << m_curSF.size() << " objects.");
        m_curSF.clear();
      }

      // Now, insert the observation in the SF:
      m_curSF.insert(obs);
    }
    else
      THROW_EXCEPTION(
          "*** ERROR *** Class is not an action or an "
          "observation");
  }
}

void RawlogGrabberApp::process_observations_for_nonsf(
    const RawlogGrabberApp::TListObservations& list_obs)
{
  // ---------------------------
  //  DO NOT USE SENSORY-FRAMES
  // ---------------------------
  for (auto& ob : list_obs)
  {
    auto& obj_ptr = ob.second;
    (*m_out_arch_ptr) << *obj_ptr;
    {
      auto lk = mrpt::lockHelper(results_mtx);
      rawlog_saved_objects++;
    }

    dump_verbose_info(obj_ptr);
  }

  if (!list_obs.empty())
  {
    MRPT_LOG_INFO_STREAM("Saved " << list_obs.size() << " objects.");
  }
}
