/* +------------------------------------------------------------------------+
   |                     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 "maps-precomp.h"  // Precomp header
//
#include <mrpt/config.h>
#include <mrpt/config/CConfigFileBase.h>
#include <mrpt/maps/CRandomFieldGridMap3D.h>
#include <mrpt/serialization/CArchive.h>
#include <mrpt/system/CTicTac.h>

#include <fstream>

using namespace mrpt;
using namespace mrpt::maps;
using namespace mrpt::system;
using namespace std;

IMPLEMENTS_SERIALIZABLE(CRandomFieldGridMap3D, CSerializable, mrpt::maps)

/*---------------------------------------------------------------
						Constructor
  ---------------------------------------------------------------*/
CRandomFieldGridMap3D::CRandomFieldGridMap3D(
	double x_min, double x_max, double y_min, double y_max, double z_min,
	double z_max, double voxel_size, bool call_initialize_now)
	: CDynamicGrid3D<TRandomFieldVoxel>(
		  x_min, x_max, y_min, y_max, z_min, z_max, voxel_size /*xy*/,
		  voxel_size /*z*/),
	  COutputLogger("CRandomFieldGridMap3D")
{
	if (call_initialize_now) this->internal_initialize();
}

/** Changes the size of the grid, erasing previous contents */
void CRandomFieldGridMap3D::setSize(
	const double x_min, const double x_max, const double y_min,
	const double y_max, const double z_min, const double z_max,
	const double resolution_xy, const double resolution_z,
	const TRandomFieldVoxel* fill_value)
{
	MRPT_START

	CDynamicGrid3D<TRandomFieldVoxel>::setSize(
		x_min, x_max, y_min, y_max, z_min, z_max, resolution_xy, resolution_z,
		fill_value);
	this->internal_initialize();

	MRPT_END
}

void CRandomFieldGridMap3D::resize(
	double new_x_min, double new_x_max, double new_y_min, double new_y_max,
	double new_z_min, double new_z_max,
	const TRandomFieldVoxel& defaultValueNewCells,
	double additionalMarginMeters)
{
	MRPT_START

	CDynamicGrid3D<TRandomFieldVoxel>::resize(
		new_x_min, new_x_max, new_y_min, new_y_max, new_z_min, new_z_max,
		defaultValueNewCells, additionalMarginMeters);
	this->internal_initialize(false);

	MRPT_END
}

void CRandomFieldGridMap3D::clear()
{
	mrpt::containers::CDynamicGrid3D<TRandomFieldVoxel>::clear();
	internal_initialize();
}

void CRandomFieldGridMap3D::internal_initialize(bool erase_prev_contents)
{
	if (erase_prev_contents)
	{
		// Set the gridmap (m_map) to initial values:
		TRandomFieldVoxel def(0, 0);  // mean, std
		fill(def);
	}

	// Reset gmrf:
	m_gmrf.setVerbosityLevel(this->getMinLoggingLevel());
	if (erase_prev_contents)
	{
		m_gmrf.clear();
		m_mrf_factors_activeObs.clear();
	}
	else
	{
		// Only remove priors, leave observations:
		m_gmrf.clearAllConstraintsByType_Binary();
	}
	m_mrf_factors_priors.clear();

	// Initiating prior:
	const size_t nodeCount = m_map.size();
	ASSERT_EQUAL_(nodeCount, m_size_x * m_size_y * m_size_z);
	ASSERT_EQUAL_(m_size_x_times_y, m_size_x * m_size_y);

	MRPT_LOG_DEBUG_STREAM(
		"[internal_initialize] Creating priors for GMRF with "
		<< nodeCount << " nodes." << std::endl);
	CTicTac tictac;
	tictac.Tic();

	m_mrf_factors_activeObs.resize(nodeCount);	// Alloc space for obs
	m_gmrf.initialize(nodeCount);

	ConnectivityDescriptor* custom_connectivity =
		m_gmrf_connectivity
			.get();	 // Use a raw ptr to avoid the cost in the inner loops

	size_t cx = 0, cy = 0, cz = 0;
	for (size_t j = 0; j < nodeCount; j++)
	{
		// add factors between this node and:
		// 1) the right node: j +1
		// 2) the back node:  j+m_size_x
		// 3) the top node: j+m_size_x*m_size*y
		//-------------------------------------------------
		const size_t c_idx_to_check[3] = {cx, cy, cz};
		const size_t c_idx_to_check_limits[3] = {
			m_size_x - 1, m_size_y - 1, m_size_z - 1};
		const size_t c_neighbor_idx_incr[3] = {1, m_size_x, m_size_x_times_y};

		for (int dir = 0; dir < 3; dir++)
		{
			if (c_idx_to_check[dir] >= c_idx_to_check_limits[dir]) continue;

			const size_t i = j + c_neighbor_idx_incr[dir];
			ASSERT_(i < nodeCount);

			double edge_lamdba = .0;
			if (custom_connectivity != nullptr)
			{
				const bool is_connected =
					custom_connectivity->getEdgeInformation(
						this, cx, cy, cz, cx + (dir == 0 ? 1 : 0),
						cy + (dir == 1 ? 1 : 0), cz + (dir == 2 ? 1 : 0),
						edge_lamdba);
				if (!is_connected) continue;
			}
			else
			{
				edge_lamdba = insertionOptions.GMRF_lambdaPrior;
			}

			TPriorFactorGMRF new_prior(*this);
			new_prior.node_id_i = i;
			new_prior.node_id_j = j;
			new_prior.Lambda = edge_lamdba;

			m_mrf_factors_priors.push_back(new_prior);
			m_gmrf.addConstraint(
				*m_mrf_factors_priors.rbegin());  // add to graph
		}

		// Increment coordinates:
		if (++cx >= m_size_x)
		{
			cx = 0;
			if (++cy >= m_size_y)
			{
				cy = 0;
				cz++;
			}
		}
	}  // end for "j"

	MRPT_LOG_DEBUG_STREAM(
		"[internal_initialize] Prior built in " << tictac.Tac() << " s\n"
												<< std::endl);
}

/*---------------------------------------------------------------
					TInsertionOptions
 ---------------------------------------------------------------*/
CRandomFieldGridMap3D::TInsertionOptions::TInsertionOptions()

	= default;

void CRandomFieldGridMap3D::TInsertionOptions::dumpToTextStream(
	std::ostream& out) const
{
	out << mrpt::format(
		"GMRF_lambdaPrior                     = %f\n", GMRF_lambdaPrior);
	out << mrpt::format(
		"GMRF_skip_variance                   = %s\n",
		GMRF_skip_variance ? "true" : "false");
}

void CRandomFieldGridMap3D::TInsertionOptions::loadFromConfigFile(
	const mrpt::config::CConfigFileBase& iniFile, const std::string& section)
{
	GMRF_lambdaPrior = iniFile.read_double(
		section.c_str(), "GMRF_lambdaPrior", GMRF_lambdaPrior);
	GMRF_skip_variance = iniFile.read_bool(
		section.c_str(), "GMRF_skip_variance", GMRF_skip_variance);
}

bool mrpt::maps::CRandomFieldGridMap3D::saveAsCSV(
	const std::string& filName_mean, const std::string& filName_stddev) const
{
	std::ofstream f_mean, f_stddev;

	f_mean.open(filName_mean);
	if (!f_mean.is_open()) { return false; }
	else
	{
		f_mean << "x coord, y coord, z coord, scalar\n";
	}

	if (!filName_stddev.empty())
	{
		f_stddev.open(filName_stddev);
		if (!f_stddev.is_open()) { return false; }
		else
		{
			f_mean << "x coord, y coord, z coord, scalar\n";
		}
	}

	const size_t nodeCount = m_map.size();
	size_t cx = 0, cy = 0, cz = 0;
	for (size_t j = 0; j < nodeCount; j++)
	{
		const double x = idx2x(cx), y = idx2y(cy), z = idx2z(cz);
		const double mean_val = m_map[j].mean_value;
		const double stddev_val = m_map[j].stddev_value;

		f_mean << mrpt::format("%f, %f, %f, %e\n", x, y, z, mean_val);

		if (f_stddev.is_open())
			f_stddev << mrpt::format("%f, %f, %f, %e\n", x, y, z, stddev_val);

		// Increment coordinates:
		if (++cx >= m_size_x)
		{
			cx = 0;
			if (++cy >= m_size_y)
			{
				cy = 0;
				cz++;
			}
		}
	}  // end for "j"

	return true;
}

void CRandomFieldGridMap3D::updateMapEstimation()
{
	ASSERTMSG_(
		!m_mrf_factors_activeObs.empty(),
		"Cannot update a map with no observations!");

	mrpt::math::CVectorDouble x_incr, x_var;
	m_gmrf.updateEstimation(
		x_incr, insertionOptions.GMRF_skip_variance ? nullptr : &x_var);

	ASSERT_(size_t(m_map.size()) == size_t(x_incr.size()));
	ASSERT_(
		insertionOptions.GMRF_skip_variance ||
		size_t(m_map.size()) == size_t(x_var.size()));

	// Update Mean-Variance in the base grid class
	for (size_t j = 0; j < m_map.size(); j++)
	{
		m_map[j].mean_value += x_incr[j];
		m_map[j].stddev_value =
			insertionOptions.GMRF_skip_variance ? .0 : std::sqrt(x_var[j]);
	}
}

void mrpt::maps::CRandomFieldGridMap3D::setVoxelsConnectivity(
	const ConnectivityDescriptor::Ptr& new_connectivity_descriptor)
{
	m_gmrf_connectivity = new_connectivity_descriptor;
}

bool CRandomFieldGridMap3D::insertIndividualReading(
	/** [in] The value observed in the (x,y,z) position */
	const double sensorReading,
	/** [in] The variance of the sensor observation */
	const double sensorVariance,
	/** [in] The (x,y,z) location */
	const mrpt::math::TPoint3D& point,
	/** [in] Voxel interpolation method: how many voxels will be affected by the
	   reading */
	const TVoxelInterpolationMethod method,
	/** [in] Run a global map update after inserting this observatin
	   (algorithm-dependant) */
	const bool update_map)
{
	MRPT_START

	ASSERT_GT_(sensorVariance, .0);
	ASSERTMSG_(
		m_mrf_factors_activeObs.size() == m_map.size(),
		"Trying to insert observation in uninitialized map (!)");

	const size_t cell_idx =
		cellAbsIndexFromCXCYCZ(x2idx(point.x), y2idx(point.y), z2idx(point.z));
	if (cell_idx == INVALID_VOXEL_IDX) return false;

	TObservationGMRF new_obs(*this);
	new_obs.node_id = cell_idx;
	new_obs.obsValue = sensorReading;
	new_obs.Lambda = 1.0 / sensorVariance;

	m_mrf_factors_activeObs[cell_idx].push_back(new_obs);
	m_gmrf.addConstraint(
		*m_mrf_factors_activeObs[cell_idx].rbegin());  // add to graph

	if (update_map) this->updateMapEstimation();

	return true;

	MRPT_END
}

uint8_t CRandomFieldGridMap3D::serializeGetVersion() const { return 0; }
void CRandomFieldGridMap3D::serializeTo(
	mrpt::serialization::CArchive& out) const
{
	dyngridcommon_writeToStream(out);

	// To assure compatibility: The size of each cell:
	auto n = static_cast<uint32_t>(sizeof(TRandomFieldVoxel));
	out << n;

	// Save the map contents:
	n = static_cast<uint32_t>(m_map.size());
	out << n;

// Save the "m_map": This requires special handling for big endian systems:
#if MRPT_IS_BIG_ENDIAN
	for (uint32_t i = 0; i < n; i++)
		out << m_map[i].mean_value << m_map[i].stddev_value;
#else
	// Little endian: just write all at once:
	out.WriteBuffer(&m_map[0], sizeof(m_map[0]) * m_map.size());
#endif

	out << insertionOptions.GMRF_lambdaPrior
		<< insertionOptions.GMRF_skip_variance;
}

void CRandomFieldGridMap3D::serializeFrom(
	mrpt::serialization::CArchive& in, uint8_t version)
{
	switch (version)
	{
		case 0:
		{
			dyngridcommon_readFromStream(in);

			// To assure compatibility: The size of each cell:
			uint32_t n;
			in >> n;

			ASSERT_EQUAL_(n, static_cast<uint32_t>(sizeof(TRandomFieldVoxel)));
			// Load the map contents:
			in >> n;
			m_map.resize(n);

// Read the note in writeToStream()
#if MRPT_IS_BIG_ENDIAN
			for (uint32_t i = 0; i < n; i++)
				in >> m_map[i].mean_value >> m_map[i].stddev_value;
#else
			// Little endian: just read all at once:
			in.ReadBuffer(&m_map[0], sizeof(m_map[0]) * m_map.size());
#endif
			in >> insertionOptions.GMRF_lambdaPrior >>
				insertionOptions.GMRF_skip_variance;
		}
		break;
		default: MRPT_THROW_UNKNOWN_SERIALIZATION_VERSION(version);
	};
}

// ============ TObservationGMRF ===========
double CRandomFieldGridMap3D::TObservationGMRF::evaluateResidual() const
{
	return m_parent->m_map[this->node_id].mean_value - this->obsValue;
}
double CRandomFieldGridMap3D::TObservationGMRF::getInformation() const
{
	return this->Lambda;
}
void CRandomFieldGridMap3D::TObservationGMRF::evalJacobian(double& dr_dx) const
{
	dr_dx = 1.0;
}
// ============ TPriorFactorGMRF ===========
double CRandomFieldGridMap3D::TPriorFactorGMRF::evaluateResidual() const
{
	return m_parent->m_map[this->node_id_i].mean_value -
		m_parent->m_map[this->node_id_j].mean_value;
}
double CRandomFieldGridMap3D::TPriorFactorGMRF::getInformation() const
{
	return this->Lambda;
}
void CRandomFieldGridMap3D::TPriorFactorGMRF::evalJacobian(
	double& dr_dx_i, double& dr_dx_j) const
{
	dr_dx_i = +1.0;
	dr_dx_j = -1.0;
}
