/* +------------------------------------------------------------------------+
   |                     Mobile Robot Programming Toolkit (MRPT)            |
   |                          https://www.mrpt.org/                         |
   |                                                                        |
   | Copyright (c) 2005-2023, 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/config/CConfigFileBase.h>
#include <mrpt/core/lock_helper.h>
#include <mrpt/img/CEnhancedMetaFile.h>
#include <mrpt/maps/COccupancyGridMap2D.h>
#include <mrpt/maps/CSimplePointsMap.h>
#include <mrpt/obs/CObservationOdometry.h>
#include <mrpt/poses/CPose3DPDFGaussian.h>
#include <mrpt/poses/CPosePDFGaussian.h>
#include <mrpt/slam/CICP.h>
#include <mrpt/slam/CMetricMapBuilderICP.h>
#include <mrpt/system/CTicTac.h>

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

CMetricMapBuilderICP::CMetricMapBuilderICP()
	: ICP_options(m_min_verbosity_level)
{
	this->setLoggerName("CMetricMapBuilderICP");
	this->initialize(CSimpleMap());
}

/*---------------------------------------------------------------
							Destructor
  ---------------------------------------------------------------*/
CMetricMapBuilderICP::~CMetricMapBuilderICP()
{
	// Ensure, we have exit all critical zones:
	enterCriticalSection();
	leaveCriticalSection();

	// Save current map to current file:
	setCurrentMapFile("");
}

/*---------------------------------------------------------------
							Options
  ---------------------------------------------------------------*/
CMetricMapBuilderICP::TConfigParams::TConfigParams(
	mrpt::system::VerbosityLevel& parent_verbosity_level)
	: matchAgainstTheGrid(false),
	  insertionLinDistance(1.0),
	  insertionAngDistance(30.0_deg),
	  localizationLinDistance(0.20),
	  localizationAngDistance(30.0_deg),
	  minICPgoodnessToAccept(0.40),
	  verbosity_level(parent_verbosity_level),
	  mapInitializers()
{
}

CMetricMapBuilderICP::TConfigParams&
	CMetricMapBuilderICP::TConfigParams::operator=(
		const CMetricMapBuilderICP::TConfigParams& other)
{
	matchAgainstTheGrid = other.matchAgainstTheGrid;
	insertionLinDistance = other.insertionLinDistance;
	insertionAngDistance = other.insertionAngDistance;
	localizationLinDistance = other.localizationLinDistance;
	localizationAngDistance = other.localizationAngDistance;
	minICPgoodnessToAccept = other.minICPgoodnessToAccept;
	//	We can't copy a reference type
	//	verbosity_level         = other.verbosity_level;
	mapInitializers = other.mapInitializers;
	return *this;
}

void CMetricMapBuilderICP::TConfigParams::loadFromConfigFile(
	const mrpt::config::CConfigFileBase& source, const std::string& section)
{
	MRPT_LOAD_CONFIG_VAR(matchAgainstTheGrid, bool, source, section)
	MRPT_LOAD_CONFIG_VAR(insertionLinDistance, double, source, section)
	MRPT_LOAD_CONFIG_VAR_DEGREES(insertionAngDistance, source, section)
	MRPT_LOAD_CONFIG_VAR(localizationLinDistance, double, source, section)
	MRPT_LOAD_CONFIG_VAR_DEGREES(localizationAngDistance, source, section)
	verbosity_level = source.read_enum<mrpt::system::VerbosityLevel>(
		section, "verbosity_level", verbosity_level);

	MRPT_LOAD_CONFIG_VAR(minICPgoodnessToAccept, double, source, section)

	mapInitializers.loadFromConfigFile(source, section);
}

void CMetricMapBuilderICP::TConfigParams::dumpToTextStream(
	std::ostream& out) const
{
	out << "\n----------- [CMetricMapBuilderICP::TConfigParams] ------------ "
		   "\n\n";

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

	out << "  Now showing 'mapsInitializers':\n";
	mapInitializers.dumpToTextStream(out);
}

/*---------------------------------------------------------------
						processObservation
 This is the new entry point of the algorithm (the old one
  was processActionObservation, which now is a wrapper to
  this method).
  ---------------------------------------------------------------*/
void CMetricMapBuilderICP::processObservation(const CObservation::Ptr& obs)
{
	auto lck = mrpt::lockHelper(critZoneChangingMap);

	MRPT_START

	if (metricMap.countMapsByClass<mrpt::maps::CPointsMap>() == 0 &&
		metricMap.countMapsByClass<mrpt::maps::COccupancyGridMap2D>() == 0)
		throw std::runtime_error(
			"Neither grid maps nor points map: Have you called initialize() "
			"after setting ICP_options.mapInitializers?");

	ASSERT_(obs);

	// Is it an odometry observation??
	if (IS_CLASS(*obs, CObservationOdometry))
	{
		MRPT_LOG_DEBUG("processObservation(): obs is CObservationOdometry");
		m_there_has_been_an_odometry = true;

		const CObservationOdometry::Ptr odo =
			std::dynamic_pointer_cast<CObservationOdometry>(obs);
		ASSERT_(odo->timestamp != INVALID_TIMESTAMP);

		CPose2D pose_before;
		bool pose_before_valid = m_lastPoseEst.getLatestRobotPose(pose_before);

		// Move our estimation:
		m_lastPoseEst.processUpdateNewOdometry(
			odo->odometry.asTPose(), odo->timestamp, odo->hasVelocities,
			odo->velocityLocal);

		if (pose_before_valid)
		{
			// Accumulate movement:
			CPose2D pose_after;
			if (m_lastPoseEst.getLatestRobotPose(pose_after))
				this->accumulateRobotDisplacementCounters(pose_after);
			MRPT_LOG_DEBUG_STREAM(
				"processObservation(): obs is CObservationOdometry, new "
				"post_after="
				<< pose_after);
		}
	}  // end it's odometry
	else
	{
		// Current robot pose given the timestamp of the observation (this can
		// include a small extrapolation
		//  using the latest known robot velocities):
		TPose2D initialEstimatedRobotPose(0, 0, 0);
		{
			mrpt::math::TTwist2D robotVelLocal, robotVelGlobal;
			if (obs->timestamp != INVALID_TIMESTAMP)
			{
				MRPT_LOG_DEBUG(
					"processObservation(): extrapolating pose from latest pose "
					"and new observation timestamp...");
				if (!m_lastPoseEst.getCurrentEstimate(
						initialEstimatedRobotPose, robotVelLocal,
						robotVelGlobal, obs->timestamp))
				{  // couldn't had a good extrapolation estimate... we'll have
					// to live with the latest pose:
					m_lastPoseEst.getLatestRobotPose(initialEstimatedRobotPose);
					MRPT_LOG_THROTTLE_WARN(
						10.0 /*seconds*/,
						"processObservation(): new pose extrapolation failed, "
						"using last pose as is.");
				}
			}
			else
			{
				MRPT_LOG_WARN(
					"processObservation(): invalid observation timestamp.");
				m_lastPoseEst.getLatestRobotPose(initialEstimatedRobotPose);
			}
		}

		// To know the total path length:
		CPose2D previousKnownRobotPose;
		m_lastPoseEst.getLatestRobotPose(previousKnownRobotPose);

		// Increment (this may only include the effects of extrapolation with
		// velocity...):
		this->accumulateRobotDisplacementCounters(
			previousKnownRobotPose);  // initialEstimatedRobotPose-previousKnownRobotPose);

		// We'll skip ICP-based localization for this observation only if:
		//  - We had some odometry since the last pose correction
		//  (m_there_has_been_an_odometry=true).
		//  - AND, the traversed distance is small enough:
		const bool we_skip_ICP_pose_correction = m_there_has_been_an_odometry &&
			m_distSinceLastICP.lin < std::min(
										 ICP_options.localizationLinDistance,
										 ICP_options.insertionLinDistance) &&
			m_distSinceLastICP.ang < std::min(
										 ICP_options.localizationAngDistance,
										 ICP_options.insertionAngDistance);

		MRPT_LOG_DEBUG_STREAM(
			"processObservation(): skipping ICP pose correction due to small "
			"odometric displacement? : "
			<< (we_skip_ICP_pose_correction ? "YES" : "NO"));

		CICP::TReturnInfo icpReturn;
		bool can_do_icp = false;

		// Select the map to match with ....
		CMetricMap* matchWith = nullptr;
		if (auto pGrid = metricMap.mapByClass<COccupancyGridMap2D>();
			ICP_options.matchAgainstTheGrid && pGrid)
		{
			matchWith = static_cast<CMetricMap*>(pGrid.get());
			MRPT_LOG_DEBUG("processObservation(): matching against gridmap.");
		}
		else
		{
			auto pPts = metricMap.mapByClass<CPointsMap>();
			ASSERTMSG_(pPts, "No points map in multi-metric map.");

			matchWith = static_cast<CMetricMap*>(pPts.get());
			MRPT_LOG_DEBUG("processObservation(): matching against point map.");
		}
		ASSERT_(matchWith != nullptr);

		if (!we_skip_ICP_pose_correction)
		{
			m_there_has_been_an_odometry = false;

			// --------------------------------------------------------------------------------------
			// Any other observation:
			//  1) If the observation generates points in a point map,
			//  do ICP 2) In any case, insert the observation if the
			//  minimum distance has been satisfaced.
			// --------------------------------------------------------------------------------------
			CSimplePointsMap sensedPoints;
			sensedPoints.insertionOptions.minDistBetweenLaserPoints = 0.02f;
			sensedPoints.insertionOptions.also_interpolate = false;

			// Create points representation of the observation:
			// Insert only those planar range scans in the altitude of
			// the grid map:
			if (ICP_options.matchAgainstTheGrid &&
				0 != metricMap.countMapsByClass<COccupancyGridMap2D>() &&
				metricMap.mapByClass<COccupancyGridMap2D>(0)
					->insertionOptions.useMapAltitude)
			{
				// Use grid altitude:
				if (IS_CLASS(*obs, CObservation2DRangeScan))
				{
					CObservation2DRangeScan::Ptr obsLaser =
						std::dynamic_pointer_cast<CObservation2DRangeScan>(obs);
					if (std::abs(
							metricMap.mapByClass<COccupancyGridMap2D>(0)
								->insertionOptions.mapAltitude -
							obsLaser->sensorPose.z()) < 0.01)
						can_do_icp = sensedPoints.insertObservationPtr(obs);
				}
			}
			else
			{
				// Do not use grid altitude:
				can_do_icp = sensedPoints.insertObservationPtr(obs);
			}

			if (IS_DERIVED(*matchWith, CPointsMap) &&
				static_cast<CPointsMap*>(matchWith)->empty())
				can_do_icp = false;	 // The reference map is empty!

			if (can_do_icp)
			{
				// We DO HAVE points with this observation:
				// Execute ICP over the current points map and the
				// sensed points:
				// ----------------------------------------------------------------------
				CICP ICP;
				ICP.options = ICP_params;

				// a first gross estimation of map 2 relative to map 1.
				const auto firstGuess =
					mrpt::poses::CPose2D(initialEstimatedRobotPose);

				CPosePDF::Ptr pestPose = ICP.Align(
					matchWith,	// Map 1
					&sensedPoints,	// Map 2
					firstGuess, icpReturn);

				if (icpReturn.goodness > ICP_options.minICPgoodnessToAccept)
				{
					// save estimation:
					CPosePDFGaussian pEst2D;
					pEst2D.copyFrom(*pestPose);

					m_lastPoseEst.processUpdateNewPoseLocalization(
						pEst2D.mean.asTPose(), obs->timestamp);
					m_lastPoseEst_cov = pEst2D.cov;

					m_distSinceLastICP.updatePose(pEst2D.mean);

					// Debug output to console:
					MRPT_LOG_INFO_STREAM(
						"processObservation: previousPose="
						<< previousKnownRobotPose << "-> currentPose="
						<< pEst2D.getMeanVal() << std::endl);
					MRPT_LOG_INFO(format(
						"[CMetricMapBuilderICP]   Fit:%.1f%% Itr:%i In "
						"%.02fms \n",
						icpReturn.goodness * 100, icpReturn.nIterations,
						1000 * icpReturn.executionTime));
				}
				else
				{
					MRPT_LOG_WARN_STREAM(
						"Ignoring ICP of low quality: " << icpReturn.goodness *
							100 << std::endl);
				}

				// Compute the transversed length:
				CPose2D currentKnownRobotPose;
				m_lastPoseEst.getLatestRobotPose(currentKnownRobotPose);

				this->accumulateRobotDisplacementCounters(
					currentKnownRobotPose);	 // currentKnownRobotPose -
				// previousKnownRobotPose);

			}  // end we can do ICP.
			else
			{
				MRPT_LOG_WARN_STREAM(
					"Cannot do ICP: empty pointmap or not suitable "
					"gridmap...\n");
			}

		}  // else, we do ICP pose correction

		// ----------------------------------------------------------
		//				CRITERION TO DECIDE MAP UPDATE:
		//   A distance large-enough from the last update for each
		//   sensor, AND
		//    either: (i) this was a good match or (ii) this is the
		//    first time for this sensor.
		// ----------------------------------------------------------
		const bool firstTimeForThisSensor =
			m_distSinceLastInsertion.find(obs->sensorLabel) ==
			m_distSinceLastInsertion.end();
		bool update = firstTimeForThisSensor ||
			((!can_do_icp ||
			  icpReturn.goodness > ICP_options.minICPgoodnessToAccept) &&
			 (m_distSinceLastInsertion[obs->sensorLabel].lin >=
				  ICP_options.insertionLinDistance ||
			  m_distSinceLastInsertion[obs->sensorLabel].ang >=
				  ICP_options.insertionAngDistance));

		// Used any "options.alwaysInsertByClass" ??
		if (options.alwaysInsertByClass.contains(obs->GetRuntimeClass()))
			update = true;

		// We need to always insert ALL the observations at the
		// beginning until the first one
		//  that actually insert some points into the map used as a
		//  reference, since otherwise we'll not be able to do ICP
		//  against an empty map!!
		if (matchWith && matchWith->isEmpty()) update = true;

		MRPT_LOG_DEBUG_STREAM(
			"update map: " << (update ? "YES" : "NO")
						   << " options.enableMapUpdating: "
						   << (options.enableMapUpdating ? "YES" : "NO"));

		if (options.enableMapUpdating && update)
		{
			CTicTac tictac;

			tictac.Tic();

			// Insert the observation:
			CPose2D currentKnownRobotPose;
			m_lastPoseEst.getLatestRobotPose(currentKnownRobotPose);

			// Create new entry:
			m_distSinceLastInsertion[obs->sensorLabel].last_update =
				currentKnownRobotPose.asTPose();

			// Reset distance counters:
			resetRobotDisplacementCounters(currentKnownRobotPose);
			// m_distSinceLastInsertion[obs->sensorLabel].updatePose(currentKnownRobotPose);

			MRPT_LOG_INFO(mrpt::format(
				"Updating map from pose %s\n",
				currentKnownRobotPose.asString().c_str()));

			CPose3D estimatedPose3D(currentKnownRobotPose);
			const bool anymap_update =
				metricMap.insertObservationPtr(obs, estimatedPose3D);
			if (!anymap_update)
				MRPT_LOG_WARN_STREAM(
					"**No map was updated** after inserting an "
					"observation of "
					"type `"
					<< obs->GetRuntimeClass()->className << "`");

			// Add to the vector of "poses"-"SFs" pairs:
			CPosePDFGaussian posePDF(currentKnownRobotPose);
			CPose3DPDF::Ptr pose3D =
				CPose3DPDF::Ptr(CPose3DPDF::createFrom2D(posePDF));

			CSensoryFrame::Ptr sf = std::make_shared<CSensoryFrame>();
			sf->insert(obs);

			SF_Poses_seq.insert(pose3D, sf);

			MRPT_LOG_INFO_STREAM(
				"Map updated OK. Done in "
				<< mrpt::system::formatTimeInterval(tictac.Tac()) << std::endl);
		}

	}  // end other observation

	// Robot path history:
	{
		TPose2D p;
		if (m_lastPoseEst.getLatestRobotPose(p)) m_estRobotPath.push_back(p);
	}

	MRPT_END

}  // end processObservation

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

						processActionObservation

  ---------------------------------------------------------------*/
void CMetricMapBuilderICP::processActionObservation(
	CActionCollection& action, CSensoryFrame& in_SF)
{
	// 1) process action:
	CActionRobotMovement2D::Ptr movEstimation =
		action.getBestMovementEstimation();
	if (movEstimation)
	{
		m_auxAccumOdometry.composeFrom(
			m_auxAccumOdometry, movEstimation->poseChange->getMeanVal());

		CObservationOdometry::Ptr obs =
			std::make_shared<CObservationOdometry>();
		obs->timestamp = movEstimation->timestamp;
		obs->odometry = m_auxAccumOdometry;
		this->processObservation(obs);
	}

	// 2) Process observations one by one:
	for (auto& i : in_SF)
		this->processObservation(i);
}

/*---------------------------------------------------------------
						setCurrentMapFile
  ---------------------------------------------------------------*/
void CMetricMapBuilderICP::setCurrentMapFile(const char* mapFile)
{
	// Save current map to current file:
	if (currentMapFile.size()) saveCurrentMapToFile(currentMapFile.c_str());

	// Sets new current map file:
	currentMapFile = mapFile;

	// Load map from file or create an empty one:
	if (currentMapFile.size()) loadCurrentMapFromFile(mapFile);
}

/*---------------------------------------------------------------
						getCurrentPoseEstimation
  ---------------------------------------------------------------*/
CPose3DPDF::Ptr CMetricMapBuilderICP::getCurrentPoseEstimation() const
{
	CPosePDFGaussian pdf2D;
	m_lastPoseEst.getLatestRobotPose(pdf2D.mean);
	pdf2D.cov = m_lastPoseEst_cov;

	CPose3DPDFGaussian::Ptr pdf3D = std::make_shared<CPose3DPDFGaussian>();
	pdf3D->copyFrom(pdf2D);
	return pdf3D;
}

/*---------------------------------------------------------------
						initialize
  ---------------------------------------------------------------*/
void CMetricMapBuilderICP::initialize(
	const CSimpleMap& initialMap, const CPosePDF* x0)
{
	MRPT_START

	// Reset vars:
	m_estRobotPath.clear();
	m_auxAccumOdometry = CPose2D(0, 0, 0);

	m_distSinceLastICP.lin = m_distSinceLastICP.ang = 0;
	m_distSinceLastInsertion.clear();

	m_there_has_been_an_odometry = false;

	// Init path & map:
	auto lck = mrpt::lockHelper(critZoneChangingMap);

	// Create metric maps:
	metricMap.setListOfMaps(ICP_options.mapInitializers);

	// copy map:
	SF_Poses_seq = initialMap;

	// Load estimated pose from given PDF:
	m_lastPoseEst.reset();

	if (x0)
		m_lastPoseEst.processUpdateNewPoseLocalization(
			x0->getMeanVal().asTPose(), mrpt::Clock::now());

	for (size_t i = 0; i < SF_Poses_seq.size(); i++)
	{
		CPose3DPDF::Ptr posePDF;
		CSensoryFrame::Ptr SF;

		// Get the SF and its pose:
		SF_Poses_seq.get(i, posePDF, SF);

		CPose3D estimatedPose3D;
		posePDF->getMean(estimatedPose3D);

		// Insert observations into the map:
		SF->insertObservationsInto(metricMap, estimatedPose3D);
	}

	MRPT_END
}

void CMetricMapBuilderICP::getCurrentMapPoints(
	std::vector<float>& x, std::vector<float>& y)
{
	auto lck = mrpt::lockHelper(critZoneChangingMap);

	auto pPts = metricMap.mapByClass<CPointsMap>(0);

	ASSERT_(pPts);
	pPts->getAllPoints(x, y);
}

/*---------------------------------------------------------------
				getCurrentlyBuiltMap
  ---------------------------------------------------------------*/
void CMetricMapBuilderICP::getCurrentlyBuiltMap(CSimpleMap& out_map) const
{
	out_map = SF_Poses_seq;
}

const CMultiMetricMap& CMetricMapBuilderICP::getCurrentlyBuiltMetricMap() const
{
	return metricMap;
}

/*---------------------------------------------------------------
			getCurrentlyBuiltMapSize
  ---------------------------------------------------------------*/
unsigned int CMetricMapBuilderICP::getCurrentlyBuiltMapSize()
{
	return SF_Poses_seq.size();
}

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

	CImage img;
	const size_t nPoses = m_estRobotPath.size();

	if (!formatEMF_BMP) THROW_EXCEPTION("Not implemented yet for BMP!");

	// grid map as bitmap:
	auto pGrid = metricMap.mapByClass<COccupancyGridMap2D>();
	if (pGrid) pGrid->getAsImage(img);

	// Draw paths (using vectorial plots!) over the EMF file:
	// -------------------------------------------------
	CEnhancedMetaFile EMF(file, 1000);

	EMF.drawImage(0, 0, img);

	unsigned int imgHeight = img.getHeight();

	// Path hypothesis:
	// ----------------------------------
	int x1, x2, y1, y2;

	// First point: (0,0)
	x2 = pGrid->x2idx(0.0f);
	y2 = pGrid->y2idx(0.0f);

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

		// Coordinates -> pixels
		x2 = pGrid->x2idx(m_estRobotPath[j].x);
		y2 = pGrid->y2idx(m_estRobotPath[j].y);

		// Draw line:
		EMF.line(
			x1, imgHeight - 1 - y1, x2, imgHeight - 1 - y2, TColor::black());
	}

	MRPT_END
}

void CMetricMapBuilderICP::accumulateRobotDisplacementCounters(
	const CPose2D& new_pose)
{
	m_distSinceLastICP.updateDistances(new_pose);
	for (auto& m : m_distSinceLastInsertion)
		m.second.updateDistances(new_pose);
}

void CMetricMapBuilderICP::resetRobotDisplacementCounters(
	const CPose2D& new_pose)
{
	m_distSinceLastICP.updatePose(new_pose);
	for (auto& m : m_distSinceLastInsertion)
		m.second.updatePose(new_pose);
}

void CMetricMapBuilderICP::TDist::updateDistances(const mrpt::poses::CPose2D& p)
{
	const auto Ap = p - mrpt::poses::CPose2D(this->last_update);
	lin = Ap.norm();
	ang = std::abs(Ap.phi());
}

void CMetricMapBuilderICP::TDist::updatePose(const mrpt::poses::CPose2D& p)
{
	this->last_update = p.asTPose();
	lin = 0;
	ang = 0;
}
