/* +------------------------------------------------------------------------+
   |                     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/io/CFileStream.h>
#include <mrpt/maps/CLandmarksMap.h>
#include <mrpt/maps/CMultiMetricMapPDF.h>
#include <mrpt/maps/COccupancyGridMap2D.h>
#include <mrpt/maps/CSimplePointsMap.h>
#include <mrpt/obs/CActionCollection.h>
#include <mrpt/obs/CActionRobotMovement2D.h>
#include <mrpt/obs/CActionRobotMovement3D.h>
#include <mrpt/obs/CObservationBeaconRanges.h>
#include <mrpt/poses/CPosePDFGaussian.h>
#include <mrpt/poses/CPosePDFGrid.h>
#include <mrpt/random.h>
#include <mrpt/serialization/CArchive.h>
#include <mrpt/system/CTicTac.h>
#include <mrpt/system/os.h>

#include "PF_aux_structs.h"

using namespace mrpt;
using namespace mrpt::math;
using namespace mrpt::slam;
using namespace mrpt::obs;
using namespace mrpt::maps;
using namespace mrpt::poses;
using namespace mrpt::random;
using namespace mrpt::system;
using namespace std;

IMPLEMENTS_SERIALIZABLE(CMultiMetricMapPDF, CSerializable, mrpt::maps)
IMPLEMENTS_SERIALIZABLE(CRBPFParticleData, CSerializable, mrpt::maps)

/*---------------------------------------------------------------
        Constructor
  ---------------------------------------------------------------*/
CMultiMetricMapPDF::CMultiMetricMapPDF(
    const bayes::CParticleFilter::TParticleFilterOptions& opts,
    const mrpt::maps::TSetOfMetricMapInitializers& mapsInitializers,
    const TPredictionParams& predictionOptions) :
    averageMap(mapsInitializers), SFs(), SF2robotPath(), options()
{
  m_particles.resize(opts.sampleSize);
  for (auto& m_particle : m_particles)
  {
    m_particle.log_w = 0;
    m_particle.d.reset(new CRBPFParticleData(mapsInitializers));
  }

  // Initialize:
  const CPose3D nullPose(0, 0, 0);
  clear(nullPose);

  // If provided, copy the whole set of params now:
  options = predictionOptions;
}

/*---------------------------------------------------------------
            clear
  ---------------------------------------------------------------*/
void CMultiMetricMapPDF::clear(const CPose2D& initialPose)
{
  CPose3D p(initialPose);
  clear(p);
}

/*---------------------------------------------------------------
            clear
  ---------------------------------------------------------------*/
void CMultiMetricMapPDF::clear(const CPose3D& initialPose)
{
  const size_t M = m_particles.size();
  for (size_t i = 0; i < M; i++)
  {
    m_particles[i].log_w = 0;

    m_particles[i].d->mapTillNow.clear();

    m_particles[i].d->robotPath.resize(1);
    m_particles[i].d->robotPath[0] = initialPose.asTPose();
  }

  SFs.clear();
  SF2robotPath.clear();

  averageMapIsUpdated = false;
}

void CMultiMetricMapPDF::clear(
    const mrpt::maps::CSimpleMap& prevMap, const mrpt::poses::CPose3D& currentPose)
{
  const size_t nParts = m_particles.size(), nOldKeyframes = prevMap.size();
  if (nOldKeyframes == 0)
  {
    // prevMap is empty, so reset the map
    clear(currentPose);
    return;
  }
  for (size_t idxPart = 0; idxPart < nParts; idxPart++)
  {
    auto& p = m_particles[idxPart];
    p.log_w = 0;

    p.d->mapTillNow.clear();

    p.d->robotPath.resize(nOldKeyframes);
    for (size_t i = 0; i < nOldKeyframes; i++)
    {
      const auto [keyframe_pose, sfkeyframe_sf, twist] = prevMap.get(i);

      // as pose, use: if the PDF is also a PF with the same number of
      // samples, use those particles;
      // otherwise, simply use the mean for all particles as an
      // approximation (with loss of uncertainty).
      mrpt::poses::CPose3D kf_pose;
      bool kf_pose_set = false;
      if (IS_CLASS(*keyframe_pose, CPose3DPDFParticles))
      {
        const auto pdf_parts = dynamic_cast<const CPose3DPDFParticles*>(keyframe_pose.get());
        ASSERT_(pdf_parts);
        if (pdf_parts->particlesCount() == nParts)
        {
          kf_pose = CPose3D(pdf_parts->m_particles[idxPart].d);
          kf_pose_set = true;
        }
      }
      if (!kf_pose_set)
      {
        kf_pose = keyframe_pose->getMeanVal();
      }
      p.d->robotPath[i] = kf_pose.asTPose();
      for (const auto& obs : *sfkeyframe_sf) p.d->mapTillNow.insertObservation(*obs, kf_pose);
    }
  }

  SFs = prevMap;  // copy
  SF2robotPath.clear();
  SF2robotPath.reserve(nOldKeyframes);
  for (size_t i = 0; i < nOldKeyframes; i++) SF2robotPath.push_back(i);

  averageMapIsUpdated = false;
}

/*---------------------------------------------------------------
            getEstimatedPosePDF
  Returns an estimate of the pose, (the mean, or mathematical expectation of the
 PDF), computed
    as a weighted average over all m_particles.
 ---------------------------------------------------------------*/
void CMultiMetricMapPDF::getEstimatedPosePDF(CPose3DPDFParticles& out_estimation) const
{
  ASSERT_(m_particles[0].d->robotPath.size() > 0);
  getEstimatedPosePDFAtTime(m_particles[0].d->robotPath.size() - 1, out_estimation);
}

/*---------------------------------------------------------------
            getEstimatedPosePDFAtTime
 ---------------------------------------------------------------*/
void CMultiMetricMapPDF::getEstimatedPosePDFAtTime(
    size_t timeStep, CPose3DPDFParticles& out_estimation) const
{
  // CPose3D	p;
  size_t i, n = m_particles.size();

  // Delete current content of "out_estimation":
  out_estimation.clearParticles();

  // Create new m_particles:
  out_estimation.m_particles.resize(n);
  for (i = 0; i < n; i++)
  {
    out_estimation.m_particles[i].d = m_particles[i].d->robotPath[timeStep];
    out_estimation.m_particles[i].log_w = m_particles[i].log_w;
  }
}

uint8_t CRBPFParticleData::serializeGetVersion() const { return 0; }
void CRBPFParticleData::serializeTo(mrpt::serialization::CArchive&) const
{
  THROW_EXCEPTION("Shouldn't arrive here");
}
void CRBPFParticleData::serializeFrom(mrpt::serialization::CArchive&, uint8_t)
{
  THROW_EXCEPTION("Shouldn't arrive here");
}

uint8_t CMultiMetricMapPDF::serializeGetVersion() const { return 0; }
void CMultiMetricMapPDF::serializeTo(mrpt::serialization::CArchive& out) const
{
  using namespace mrpt::serialization;

  out.WriteAs<uint32_t>(m_particles.size());
  for (const auto& part : m_particles)
  {
    out << part.log_w;
    out << part.d->mapTillNow;
    out.WriteAs<uint32_t>(part.d->robotPath.size());
    for (const auto& p : part.d->robotPath) out << p;
  }
  out << SFs << SF2robotPath;
}

void CMultiMetricMapPDF::serializeFrom(mrpt::serialization::CArchive& in, uint8_t version)
{
  switch (version)
  {
    case 0:
    {
      uint32_t i, n, j, m;

      // Delete current contents:
      // --------------------------
      clearParticles();
      SFs.clear();
      SF2robotPath.clear();

      averageMapIsUpdated = false;

      // Load the new data:
      // --------------------
      in >> n;

      m_particles.resize(n);
      for (i = 0; i < n; i++)
      {
        m_particles[i].d.reset(new CRBPFParticleData());

        // Load
        in >> m_particles[i].log_w >> m_particles[i].d->mapTillNow;

        in >> m;
        m_particles[i].d->robotPath.resize(m);
        for (j = 0; j < m; j++) in >> m_particles[i].d->robotPath[j];
      }

      in >> SFs >> SF2robotPath;
    }
    break;
    default:
      MRPT_THROW_UNKNOWN_SERIALIZATION_VERSION(version);
  };
}

TPose3D CMultiMetricMapPDF::getLastPose(size_t i, bool& is_valid_pose) const
{
  if (i >= m_particles.size()) THROW_EXCEPTION("Particle index out of bounds!");

  if (m_particles[i].d->robotPath.empty())
  {
    is_valid_pose = false;
    return TPose3D(0, 0, 0, 0, 0, 0);
  }
  else
  {
    return *m_particles[i].d->robotPath.rbegin();
  }
}

const CMultiMetricMap* CMultiMetricMapPDF::getAveragedMetricMapEstimation()
{
  rebuildAverageMap();
  return &averageMap;
}

void CMultiMetricMapPDF::rebuildAverageMap()
{
  float min_x = 1e6, max_x = -1e6, min_y = 1e6, max_y = -1e6;

  if (averageMapIsUpdated) return;

  // ---------------------------------------------------------
  //					GRID
  // ---------------------------------------------------------
  for (auto& p : m_particles)
  {
    ASSERT_(p.d->mapTillNow.countMapsByClass<COccupancyGridMap2D>() > 0);

    auto grid = p.d->mapTillNow.mapByClass<COccupancyGridMap2D>(0);

    min_x = min(min_x, grid->getXMin());
    max_x = max(max_x, grid->getXMax());
    min_y = min(min_y, grid->getYMin());
    max_y = max(max_y, grid->getYMax());
  }

  // Ensure all maps have the same dimensions:
  for (auto& p : m_particles)
  {
    auto grid = p.d->mapTillNow.mapByClass<COccupancyGridMap2D>(0);
    grid->resizeGrid(min_x, max_x, min_y, max_y, 0.5f, false);
  }

  float grid_resolution = 0.1f;
  for (auto& p : m_particles)
  {
    auto grid = p.d->mapTillNow.mapByClass<COccupancyGridMap2D>(0);
    grid_resolution = grid->getResolution();

    min_x = min(min_x, grid->getXMin());
    max_x = max(max_x, grid->getXMax());
    min_y = min(min_y, grid->getYMin());
    max_y = max(max_y, grid->getYMax());
  }

  // Prepare target map:
  ASSERT_(averageMap.countMapsByClass<COccupancyGridMap2D>() > 0);
  auto avrg_grid = averageMap.mapByClass<COccupancyGridMap2D>(0);
  avrg_grid->setSize(min_x, max_x, min_y, max_y, grid_resolution, 0);

  // Compute the sum of weights:
  double sumLinearWeights = 0;
  for (auto& p : m_particles) sumLinearWeights += exp(p.log_w);

  // CHECK:
  for (auto& p : m_particles)
  {
    auto grid = p.d->mapTillNow.mapByClass<COccupancyGridMap2D>(0);
    ASSERT_(grid->getSizeX() == avrg_grid->getSizeX());
    ASSERT_(grid->getSizeY() == avrg_grid->getSizeY());
  }

  {
    // ******************************************************
    //     Implementation WITHOUT the SSE Instructions Set
    // ******************************************************
    MRPT_START

    // Reserve a float grid-map, add weight all maps
    // -------------------------------------------------------------------------------------------
    std::vector<float> floatMap;
    floatMap.resize(avrg_grid->m_map.size(), 0);

    // For each particle in the RBPF:
    double sumW = 0;
    for (auto& p : m_particles) sumW += exp(p.log_w);

    if (sumW == 0) sumW = 1;

    for (auto& p : m_particles)
    {
      auto grid = p.d->mapTillNow.mapByClass<COccupancyGridMap2D>(0);
      // Variables:
      std::vector<COccupancyGridMap2D::cellType>::iterator srcCell;
      auto firstSrcCell = grid->m_map.begin();
      auto lastSrcCell = grid->m_map.end();
      std::vector<float>::iterator destCell;

      // The weight of particle:
      float w = exp(p.log_w) / sumW;

      ASSERT_(grid->m_map.size() == floatMap.size());

      // For each cell in individual maps:
      for (srcCell = firstSrcCell, destCell = floatMap.begin(); srcCell != lastSrcCell;
           srcCell++, destCell++)
        (*destCell) += w * (*srcCell);
    }

    // Copy to fixed point map:
    std::vector<float>::iterator srcCell;
    auto destCell = avrg_grid->m_map.begin();

    ASSERT_(avrg_grid->m_map.size() == floatMap.size());

    for (srcCell = floatMap.begin(); srcCell != floatMap.end(); srcCell++, destCell++)
      *destCell = static_cast<COccupancyGridMap2D::cellType>(*srcCell);

    MRPT_END
  }  // End of SSE not supported

  // Don't calculate again until really necessary.
  averageMapIsUpdated = true;
}

/*---------------------------------------------------------------
            insertObservation
 ---------------------------------------------------------------*/
bool CMultiMetricMapPDF::insertObservation(CSensoryFrame& sf)
{
  const size_t M = particlesCount();

  // Insert into SFs:
  CPose3DPDFParticles::Ptr posePDF = std::make_shared<CPose3DPDFParticles>();
  getEstimatedPosePDF(*posePDF);

  // Insert it into the SFs and the SF2robotPath list:
  const uint32_t new_sf_id = SFs.size();
  SFs.insert(posePDF, CSensoryFrame::Create(sf));
  SF2robotPath.resize(new_sf_id + 1);
  SF2robotPath[new_sf_id] = m_particles[0].d->robotPath.size() - 1;

  bool anymap = false;
  for (size_t i = 0; i < M; i++)
  {
    bool pose_is_valid;
    const CPose3D robotPose = CPose3D(getLastPose(i, pose_is_valid));
    // ASSERT_(pose_is_valid); // if not, use the default (0,0,0)
    const bool map_modified = sf.insertObservationsInto(m_particles[i].d->mapTillNow, robotPose);
    anymap = anymap || map_modified;
  }

  averageMapIsUpdated = false;
  return anymap;
}

/*---------------------------------------------------------------
            getPath
 ---------------------------------------------------------------*/
void CMultiMetricMapPDF::getPath(size_t i, std::deque<math::TPose3D>& out_path) const
{
  if (i >= m_particles.size()) THROW_EXCEPTION("Index out of bounds");
  out_path = m_particles[i].d->robotPath;
}

/*---------------------------------------------------------------
          getCurrentEntropyOfPaths
  ---------------------------------------------------------------*/
double CMultiMetricMapPDF::getCurrentEntropyOfPaths()
{
  size_t i;
  size_t N = m_particles[0].d->robotPath.size();  // The poses count along the paths

  // Compute paths entropy:
  // ---------------------------
  double H_paths = 0;

  if (N)
  {
    // For each pose along the path:
    for (i = 0; i < N; i++)
    {
      // Get pose est. as m_particles:
      CPose3DPDFParticles posePDFParts;
      getEstimatedPosePDFAtTime(i, posePDFParts);

      // Approximate to gaussian and compute entropy of covariance:
      H_paths += posePDFParts.getCovarianceEntropy();
    }
    H_paths /= N;
  }
  return H_paths;
}

/*---------------------------------------------------------------
          getCurrentJointEntropy
  ---------------------------------------------------------------*/
double CMultiMetricMapPDF::getCurrentJointEntropy()
{
  double H_joint, H_paths, H_maps;
  COccupancyGridMap2D::TEntropyInfo entropy;

  // Entropy of the paths:
  H_paths = getCurrentEntropyOfPaths();

  float min_x = 1e6, max_x = -1e6, min_y = 1e6, max_y = -1e6;

  // Make sure all grids have the same size
  for (auto& p : m_particles)
  {
    auto grid = p.d->mapTillNow.mapByClass<COccupancyGridMap2D>(0);
    ASSERT_(grid);

    min_x = min(min_x, grid->getXMin());
    max_x = max(max_x, grid->getXMax());
    min_y = min(min_y, grid->getYMin());
    max_y = max(max_y, grid->getYMax());
  }

  // Ensure all maps have the same dimensions:
  for (auto& p : m_particles)
  {
    auto grid = p.d->mapTillNow.mapByClass<COccupancyGridMap2D>(0);
    grid->resizeGrid(min_x, max_x, min_y, max_y, 0.5f, false);
  }

  // Sum of linear weights:
  double sumLinearWeights = 0;
  for (auto& p : m_particles) sumLinearWeights += exp(p.log_w);

  // Compute weighted maps entropy:
  // --------------------------------
  H_maps = 0;
  for (auto& p : m_particles)
  {
    auto grid = p.d->mapTillNow.mapByClass<COccupancyGridMap2D>(0);

    grid->computeEntropy(entropy);
    H_maps += exp(p.log_w) * entropy.H / sumLinearWeights;
  }

  printf("H_paths=%e\n", H_paths);
  printf("H_maps=%e\n", H_maps);

  H_joint = H_paths + H_maps;
  return H_joint;
}

const CMultiMetricMap* CMultiMetricMapPDF::getCurrentMostLikelyMetricMap() const
{
  size_t i, max_i = 0, n = m_particles.size();
  double max_w = m_particles[0].log_w;

  for (i = 0; i < n; i++)
  {
    if (m_particles[i].log_w > max_w)
    {
      max_w = m_particles[i].log_w;
      max_i = i;
    }
  }

  // Return its map:
  return &m_particles[max_i].d->mapTillNow;
}

/*---------------------------------------------------------------
        updateSensoryFrameSequence
  ---------------------------------------------------------------*/
void CMultiMetricMapPDF::updateSensoryFrameSequence()
{
  MRPT_START
  CPose3DPDFParticles posePartsPDF;

  for (size_t i = 0; i < SFs.size(); i++)
  {
    // Get last estimation:
    auto& kf = SFs.get(i);

    // Compute the new one:
    getEstimatedPosePDFAtTime(SF2robotPath[i], posePartsPDF);

    // Copy into SFs:
    kf.pose->copyFrom(posePartsPDF);
  }

  MRPT_END
}

/*---------------------------------------------------------------
        saveCurrentPathEstimationToTextFile
  ---------------------------------------------------------------*/
void CMultiMetricMapPDF::saveCurrentPathEstimationToTextFile(const std::string& fil)
{
  FILE* f = os::fopen(fil.c_str(), "wt");
  if (!f) return;

  for (auto& m_particle : m_particles)
  {
    for (size_t i = 0; i < m_particle.d->robotPath.size(); i++)
    {
      const mrpt::math::TPose3D& p = m_particle.d->robotPath[i];

      os::fprintf(f, "%.04f %.04f %.04f %.04f %.04f %.04f ", p.x, p.y, p.z, p.yaw, p.pitch, p.roll);
    }
    os::fprintf(f, " %e\n", m_particle.log_w);
  }

  os::fclose(f);
}

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

  out << mrpt::format(
      "pfOptimalProposal_mapSelection          = %i\n", pfOptimalProposal_mapSelection);
  out << mrpt::format("ICPGlobalAlign_MinQuality               = %f\n", ICPGlobalAlign_MinQuality);

  KLD_params.dumpToTextStream(out);
  icp_params.dumpToTextStream(out);
  out << "\n";
}

/*---------------------------------------------------------------
          loadFromConfigFile
  ---------------------------------------------------------------*/
void CMultiMetricMapPDF::TPredictionParams::loadFromConfigFile(
    const mrpt::config::CConfigFileBase& iniFile, const std::string& section)
{
  pfOptimalProposal_mapSelection = iniFile.read_int(
      section, "pfOptimalProposal_mapSelection", pfOptimalProposal_mapSelection, true);

  MRPT_LOAD_CONFIG_VAR(ICPGlobalAlign_MinQuality, float, iniFile, section);

  KLD_params.loadFromConfigFile(iniFile, section);
  icp_params.loadFromConfigFile(iniFile, section);
}
