/* +------------------------------------------------------------------------+
   |                     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 "slam-precomp.h"  // Precompiled headers
//
#include <mrpt/core/lock_helper.h>
#include <mrpt/core/round.h>
#include <mrpt/img/CEnhancedMetaFile.h>
#include <mrpt/maps/COccupancyGridMap2D.h>
#include <mrpt/obs/CActionRobotMovement3D.h>
#include <mrpt/obs/CObservationStereoImages.h>
#include <mrpt/slam/CMetricMapBuilderRBPF.h>

using namespace mrpt;
using namespace mrpt::slam;
using namespace mrpt::math;
using namespace mrpt::maps;
using namespace mrpt::obs;
using namespace mrpt::poses;
using namespace mrpt::bayes;
using namespace mrpt::img;

/*---------------------------------------------------------------
            Constructor
  ---------------------------------------------------------------*/
CMetricMapBuilderRBPF::CMetricMapBuilderRBPF(const TConstructionOptions& initializationOptions) :
    mapPDF(
        initializationOptions.PF_options,
        initializationOptions.mapsInitializers,
        initializationOptions.predictionOptions),
    m_PF_options(initializationOptions.PF_options),
    insertionLinDistance(initializationOptions.insertionLinDistance),
    insertionAngDistance(initializationOptions.insertionAngDistance),
    localizeLinDistance(initializationOptions.localizeLinDistance),
    localizeAngDistance(initializationOptions.localizeAngDistance),
    odoIncrementSinceLastLocalization(),
    odoIncrementSinceLastMapUpdate()
{
  setLoggerName("CMetricMapBuilderRBPF");
  setVerbosityLevel(initializationOptions.verbosity_level);
  // Reset:
  clear();
}

CMetricMapBuilderRBPF::CMetricMapBuilderRBPF()
{
  this->setLoggerName("CMetricMapBuilderRBPF");
  MRPT_LOG_WARN("Empty constructor invoked!\n");
}

CMetricMapBuilderRBPF& CMetricMapBuilderRBPF::operator=(const CMetricMapBuilderRBPF& src)
{
  if (this == &src)
  {
    return *this;
  }
  mapPDF = src.mapPDF;
  m_PF_options = src.m_PF_options;
  insertionLinDistance = src.insertionLinDistance;
  insertionAngDistance = src.insertionAngDistance;
  localizeLinDistance = src.localizeLinDistance;
  localizeAngDistance = src.localizeAngDistance;
  odoIncrementSinceLastLocalization = src.odoIncrementSinceLastLocalization;
  odoIncrementSinceLastMapUpdate = src.odoIncrementSinceLastMapUpdate;
  m_statsLastIteration = src.m_statsLastIteration;
  return *this;
}

/*---------------------------------------------------------------
            Destructor
  ---------------------------------------------------------------*/
CMetricMapBuilderRBPF::~CMetricMapBuilderRBPF() = default;
/*---------------------------------------------------------------
            clear
  ---------------------------------------------------------------*/
void CMetricMapBuilderRBPF::clear()
{
  auto lck = mrpt::lockHelper(critZoneChangingMap);

  MRPT_LOG_DEBUG("CMetricMapBuilderRBPF::clear() called.");
  static CPose2D nullPose(0, 0, 0);

  // Reset traveled distances counters:
  odoIncrementSinceLastLocalization = CPose3DPDFGaussian();

  odoIncrementSinceLastMapUpdate.setFromValues(0, 0, 0, 0, 0, 0);

  // Clear maps for each particle:
  mapPDF.clear(nullPose);
}

/*---------------------------------------------------------------
          processActionObservation
  ---------------------------------------------------------------*/
void CMetricMapBuilderRBPF::processActionObservation(
    CActionCollection& action, CSensoryFrame& observations)
{
  MRPT_START
  auto lck = mrpt::lockHelper(critZoneChangingMap);

  // Update the traveled distance estimations:
  {
    CActionRobotMovement3D::Ptr act3D = action.getActionByClass<CActionRobotMovement3D>();
    CActionRobotMovement2D::Ptr act2D = action.getActionByClass<CActionRobotMovement2D>();
    if (act3D)
    {
      MRPT_LOG_DEBUG_STREAM(
          "processActionObservation(): Input action is "
          "CActionRobotMovement3D="
          << act3D->poseChange.getMeanVal().asString());
      odoIncrementSinceLastMapUpdate += act3D->poseChange.getMeanVal();
      odoIncrementSinceLastLocalization += act3D->poseChange;
    }
    else if (act2D)
    {
      MRPT_LOG_DEBUG_STREAM(
          "processActionObservation(): Input action is "
          "CActionRobotMovement2D="
          << act2D->poseChange->getMeanVal().asString());
      odoIncrementSinceLastMapUpdate += mrpt::poses::CPose3D(act2D->poseChange->getMeanVal());
      odoIncrementSinceLastLocalization.mean +=
          mrpt::poses::CPose3D(act2D->poseChange->getMeanVal());
    }
    else
    {
      MRPT_LOG_WARN("Action contains no odometry.\n");
    }
  }

  // Execute particle filter:
  //   But only if the traveled distance since the last update is long enough,
  //    or this is the first observation, etc...
  // ----------------------------------------------------------------------------
  bool do_localization =
      (!mapPDF.SFs.size() ||  // This is the first observation!
       options.debugForceInsertion ||
       odoIncrementSinceLastLocalization.mean.norm() > localizeLinDistance ||
       std::abs(odoIncrementSinceLastLocalization.mean.yaw()) > localizeAngDistance);

  bool do_map_update =
      (!mapPDF.SFs.size() ||  // This is the first observation!
       options.debugForceInsertion ||
       odoIncrementSinceLastMapUpdate.norm() > insertionLinDistance ||
       std::abs(odoIncrementSinceLastMapUpdate.yaw()) > insertionAngDistance);

  // Used any "options.alwaysInsertByClass" ??
  for (auto itCl = options.alwaysInsertByClass.data.begin();
       !do_map_update && itCl != options.alwaysInsertByClass.data.end(); ++itCl)
    for (auto& o : observations)
      if (o->GetRuntimeClass() == *itCl)
      {
        do_map_update = true;
        do_localization = true;
        break;
      }

  if (do_map_update) do_localization = true;

  MRPT_LOG_DEBUG(mrpt::format(
      "do_map_update=%s do_localization=%s", do_map_update ? "YES" : "NO",
      do_localization ? "YES" : "NO"));

  if (do_localization)
  {
    // Create an artificial action object, since
    //  we've been collecting them until a threshold:
    // ------------------------------------------------
    CActionCollection fakeActs;
    {
      CActionRobotMovement3D::Ptr act3D = action.getActionByClass<CActionRobotMovement3D>();
      if (act3D)
      {
        CActionRobotMovement3D newAct;
        newAct.estimationMethod = act3D->estimationMethod;
        newAct.poseChange = odoIncrementSinceLastLocalization;
        newAct.timestamp = act3D->timestamp;
        fakeActs.insert(newAct);
      }
      else
      {
        // It must be 2D odometry:
        CActionRobotMovement2D::Ptr act2D = action.getActionByClass<CActionRobotMovement2D>();
        ASSERT_(act2D);
        CActionRobotMovement2D newAct;
        newAct.computeFromOdometry(
            CPose2D(odoIncrementSinceLastLocalization.mean), act2D->motionModelConfiguration);
        newAct.timestamp = act2D->timestamp;
        fakeActs.insert(newAct);
      }
    }

    MRPT_LOG_DEBUG_STREAM(
        "odoIncrementSinceLastLocalization before resetting = "
        << odoIncrementSinceLastLocalization.mean);
    // Reset distance counters:
    odoIncrementSinceLastLocalization.mean.setFromValues(0, 0, 0, 0, 0, 0);
    odoIncrementSinceLastLocalization.cov.setZero();

    CParticleFilter pf;
    pf.m_options = m_PF_options;
    pf.setVerbosityLevel(this->getMinLoggingLevel());

    pf.executeOn(mapPDF, &fakeActs, &observations);

    if (isLoggingLevelVisible(mrpt::system::LVL_INFO))
    {
      // Get current pose estimation:
      CPose3DPDFParticles poseEstimation;
      CPose3D meanPose;
      mapPDF.getEstimatedPosePDF(poseEstimation);
      poseEstimation.getMean(meanPose);

      const auto [cov, estPos] = poseEstimation.getCovarianceAndMean();

      MRPT_LOG_INFO_STREAM(
          "New pose=" << estPos << std::endl
                      << "New ESS:" << mapPDF.ESS() << std::endl);
      MRPT_LOG_INFO(format(
          "   STDs: x=%2.3f y=%2.3f z=%.03f yaw=%2.3fdeg\n", sqrt(cov(0, 0)), sqrt(cov(1, 1)),
          sqrt(cov(2, 2)), RAD2DEG(sqrt(cov(3, 3)))));
    }
  }

  if (do_map_update)
  {
    odoIncrementSinceLastMapUpdate.setFromValues(0, 0, 0, 0, 0, 0);

    // Update the particles' maps:
    // -------------------------------------------------
    MRPT_LOG_INFO("New observation inserted into the map.");

    // Add current observation to the map:
    const bool anymap_update = mapPDF.insertObservation(observations);
    if (!anymap_update)
      MRPT_LOG_WARN_STREAM(
          "**No map was updated** after inserting a CSensoryFrame with " << observations.size());

    m_statsLastIteration.observationsInserted = true;
  }
  else
  {
    m_statsLastIteration.observationsInserted = false;
  }

  // Added 29/JUN/2007 JLBC: Tell all maps that they can now free aux.
  // variables
  //  (if any) since one PF cycle is over:
  for (auto& m_particle : mapPDF.m_particles) m_particle.d->mapTillNow.auxParticleFilterCleanUp();

  MRPT_END
}

/*---------------------------------------------------------------
          initialize
  ---------------------------------------------------------------*/
void CMetricMapBuilderRBPF::initialize(const CSimpleMap& initialMap, const CPosePDF* x0)
{
  MRPT_LOG_INFO_STREAM("[initialize] Called with " << initialMap.size() << " nodes in fixed map");

  this->clear();

  auto lck = mrpt::lockHelper(critZoneChangingMap);

  mrpt::poses::CPose3D curPose;
  if (x0)
  {
    curPose = mrpt::poses::CPose3D(x0->getMeanVal());
  }
  else if (!initialMap.empty())
  {
    // get pose of last keyframe:
    curPose = initialMap.rbegin()->pose->getMeanVal();
  }
  MRPT_LOG_INFO_STREAM("[initialize] Initial pose: " << curPose);

  // Clear maps for each particle & set pose:
  mapPDF.clear(initialMap, curPose);
}

/*---------------------------------------------------------------
            getCurrentPoseEstimation
  ---------------------------------------------------------------*/
CPose3DPDF::Ptr CMetricMapBuilderRBPF::getCurrentPoseEstimation() const
{
  CPose3DPDFParticles::Ptr posePDF = std::make_shared<CPose3DPDFParticles>();
  mapPDF.getEstimatedPosePDF(*posePDF);

  // Adds additional increment from accumulated odometry since last
  // localization update:
  for (auto& p : posePDF->m_particles)
  {
    p.d.composePose(this->odoIncrementSinceLastLocalization.mean.asTPose(), p.d);
  }
  return posePDF;
}

/*---------------------------------------------------------------
            getCurrentMostLikelyPath
  ---------------------------------------------------------------*/
void CMetricMapBuilderRBPF::getCurrentMostLikelyPath(std::deque<TPose3D>& outPath) const
{
  double maxW = -1, w;
  size_t mostLik = 0;
  for (size_t i = 0; i < mapPDF.particlesCount(); i++)
  {
    w = mapPDF.getW(i);
    if (w > maxW)
    {
      maxW = w;
      mostLik = i;
    }
  }

  mapPDF.getPath(mostLik, outPath);
}

/*---------------------------------------------------------------
            getCurrentlyBuiltMap
  ---------------------------------------------------------------*/
void CMetricMapBuilderRBPF::getCurrentlyBuiltMap(CSimpleMap& out_map) const
{
  const_cast<CMetricMapBuilderRBPF*>(this)->mapPDF.updateSensoryFrameSequence();
  out_map = mapPDF.SFs;
}

const CMultiMetricMap& CMetricMapBuilderRBPF::getCurrentlyBuiltMetricMap() const
{
  auto ptr = mapPDF.getCurrentMostLikelyMetricMap();
  ASSERT_(ptr);
  return *ptr;
}

/*---------------------------------------------------------------
      getCurrentlyBuiltMapSize
  ---------------------------------------------------------------*/
unsigned int CMetricMapBuilderRBPF::getCurrentlyBuiltMapSize() { return mapPDF.SFs.size(); }

void CMetricMapBuilderRBPF::drawCurrentEstimationToImage(CCanvas* img)
{
  using mrpt::round;

  unsigned int i, M = mapPDF.particlesCount();
  std::deque<TPose3D> path;
  unsigned int imgHeight = 0;

  MRPT_START

  const auto* curMap = mapPDF.getCurrentMostLikelyMetricMap();

  ASSERT_(curMap->countMapsByClass<COccupancyGridMap2D>() > 0);

  // Find which is the most likely path index:
  unsigned int bestPath = 0;
  double bestPathLik = -1;
  for (i = 0; i < M; i++)
  {
    if (mapPDF.getW(i) > bestPathLik)
    {
      bestPathLik = mapPDF.getW(i);
      bestPath = i;
    }
  }

  // Compute the length of the paths:
  mapPDF.getPath(0, path);

  // Adapt the canvas size:
  bool alreadyCopiedImage = false;
  auto g = curMap->mapByClass<COccupancyGridMap2D>(0);
  {
    auto* obj = dynamic_cast<CImage*>(img);
    if (obj) obj->resize(g->getSizeX(), g->getSizeY(), mrpt::img::CH_GRAY);
  }
  if (!alreadyCopiedImage)
  {
    CImage imgGrid;

    // grid map as bitmap:
    // ----------------------------------
    g->getAsImage(imgGrid);

    img->drawImage(0, 0, imgGrid);
    imgHeight = imgGrid.getHeight();
  }

  int x1 = 0, x2 = 0, y1 = 0, y2 = 0;
  float x_min = g->getXMin();
  float y_min = g->getYMin();
  float resolution = g->getResolution();

  // Paths hypothesis:
  // ----------------------------------
  /***/
  for (i = 0; i <= M; i++)
  {
    if (i != bestPath || i == M)
    {
      mapPDF.getPath(i == M ? bestPath : i, path);

      size_t nPoses = path.size();

      // First point: (0,0)
      x2 = round((path[0].x - x_min) / resolution);
      y2 = round((path[0].y - y_min) / resolution);

      // Draw path in the bitmap:
      for (size_t j = 0; j < nPoses; j++)
      {
        // For next segment
        x1 = x2;
        y1 = y2;

        // Coordinates -> pixels
        x2 = round((path[j].x - x_min) / resolution);
        y2 = round((path[j].y - y_min) / resolution);

        // Draw line:
        img->line(
            x1, round((imgHeight - 1) - y1), x2, round((imgHeight - 1) - y2),
            i == M ? TColor(0, 0, 0) : TColor(0x50, 0x50, 0x50),  // Color, gray levels,
            i == M ? 3 : 1                                        // Line width
        );
      }
    }
  }

  MRPT_END
}

/*---------------------------------------------------------------
          saveCurrentEstimationToImage
  ---------------------------------------------------------------*/
void CMetricMapBuilderRBPF::saveCurrentEstimationToImage(
    const std::string& file, bool formatEMF_BMP)
{
  MRPT_START

  if (formatEMF_BMP)
  {
    // Draw paths (using vectorial plots!) over the EMF file:
    // --------------------------------------------------------
    CEnhancedMetaFile EMF(file, 100 /* Scale */);
    drawCurrentEstimationToImage(&EMF);
  }
  else
  {
    CImage img(1, 1, CH_GRAY);
    drawCurrentEstimationToImage(&img);
    bool savedOk = img.saveToFile(file);
    ASSERT_(savedOk);
  }

  MRPT_END
}

/*---------------------------------------------------------------
          getCurrentJointEntropy
  ---------------------------------------------------------------*/
double CMetricMapBuilderRBPF::getCurrentJointEntropy() { return mapPDF.getCurrentJointEntropy(); }

/*---------------------------------------------------------------
          saveCurrentPathEstimationToTextFile
  ---------------------------------------------------------------*/
void CMetricMapBuilderRBPF::saveCurrentPathEstimationToTextFile(const std::string& fil)
{
  mapPDF.saveCurrentPathEstimationToTextFile(fil);
}

/*---------------------------------------------------------------
            TConstructionOptions
  ---------------------------------------------------------------*/
CMetricMapBuilderRBPF::TConstructionOptions::TConstructionOptions() :
    insertionAngDistance(30.0_deg),

    localizeAngDistance(10.0_deg),
    PF_options(),
    mapsInitializers(),
    predictionOptions()

{
}

/*---------------------------------------------------------------
          dumpToTextStream
  ---------------------------------------------------------------*/
void CMetricMapBuilderRBPF::TConstructionOptions::dumpToTextStream(std::ostream& out) const
{
  out << "\n----------- [CMetricMapBuilderRBPF::TConstructionOptions] "
         "------------ \n\n";

  out << mrpt::format("insertionLinDistance                    = %f m\n", insertionLinDistance);
  out << mrpt::format(
      "insertionAngDistance                    = %f deg\n", RAD2DEG(insertionAngDistance));
  out << mrpt::format("localizeLinDistance                     = %f m\n", localizeLinDistance);
  out << mrpt::format(
      "localizeAngDistance                     = %f deg\n", RAD2DEG(localizeAngDistance));
  out << mrpt::format(
      "verbosity_level                         = %s\n",
      mrpt::typemeta::TEnumType<mrpt::system::VerbosityLevel>::value2name(verbosity_level).c_str());

  PF_options.dumpToTextStream(out);

  out << "  Now showing 'mapsInitializers' and 'predictionOptions':\n";
  out << "\n";

  mapsInitializers.dumpToTextStream(out);
  predictionOptions.dumpToTextStream(out);
}

/*---------------------------------------------------------------
          loadFromConfigFile
  ---------------------------------------------------------------*/
void CMetricMapBuilderRBPF::TConstructionOptions::loadFromConfigFile(
    const mrpt::config::CConfigFileBase& iniFile, const std::string& section)
{
  MRPT_START

  PF_options.loadFromConfigFile(iniFile, section);

  MRPT_LOAD_CONFIG_VAR(insertionLinDistance, float, iniFile, section);
  MRPT_LOAD_HERE_CONFIG_VAR_DEGREES_NO_DEFAULT(
      insertionAngDistance_deg, double, insertionAngDistance, iniFile, section);

  MRPT_LOAD_CONFIG_VAR(localizeLinDistance, float, iniFile, section);
  MRPT_LOAD_HERE_CONFIG_VAR_DEGREES_NO_DEFAULT(
      localizeAngDistance_deg, double, localizeAngDistance, iniFile, section);
  verbosity_level =
      iniFile.read_enum<mrpt::system::VerbosityLevel>(section, "verbosity_level", verbosity_level);

  mapsInitializers.loadFromConfigFile(iniFile, section);
  predictionOptions.loadFromConfigFile(iniFile, section);

  MRPT_END
}
