/* +------------------------------------------------------------------------+
   |                     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 "poses-precomp.h"	// Precompiled headers
//
#include <mrpt/math/distributions.h>
#include <mrpt/math/wrap2pi.h>
#include <mrpt/poses/CPosePDFGaussian.h>
#include <mrpt/poses/CPosePDFParticles.h>
#include <mrpt/poses/SO_SE_average.h>
#include <mrpt/random.h>
#include <mrpt/serialization/CArchive.h>
#include <mrpt/system/os.h>

#include <fstream>

using namespace mrpt;
using namespace mrpt::bayes;
using namespace mrpt::poses;
using namespace mrpt::math;
using namespace mrpt::random;
using namespace mrpt::system;
using namespace std;

IMPLEMENTS_SERIALIZABLE(CPosePDFParticles, CPosePDF, mrpt::poses)

CPosePDFParticles::CPosePDFParticles(size_t M)
{
	m_particles.resize(M);
	for (auto& p : m_particles)
	{
		p.log_w = .0;
		p.d = TPose2D();
	}
	TPose2D nullPose(0, 0, 0);
	resetDeterministic(nullPose);
}

void CPosePDFParticles::copyFrom(const CPosePDF& o)
{
	MRPT_START

	CParticleList::iterator itDest;
	CParticleList::const_iterator itSrc;

	if (this == &o) return;	 // It may be used sometimes

	if (o.GetRuntimeClass() == CLASS_ID(CPosePDFParticles))
	{
		const auto* pdf = dynamic_cast<const CPosePDFParticles*>(&o);
		ASSERT_(pdf);

		// Both are m_particles:
		m_particles = pdf->m_particles;
	}
	else if (o.GetRuntimeClass() == CLASS_ID(CPosePDFGaussian))
	{
		const auto* pdf = dynamic_cast<const CPosePDFGaussian*>(&o);
		size_t M = m_particles.size();
		std::vector<CVectorDouble> parts;
		std::vector<CVectorDouble>::iterator partsIt;

		getRandomGenerator().drawGaussianMultivariateMany(parts, M, pdf->cov);

		clearParticles();
		m_particles.resize(M);

		for (itDest = m_particles.begin(), partsIt = parts.begin();
			 itDest != m_particles.end(); ++itDest, ++partsIt)
		{
			itDest->log_w = 0;
			itDest->d = TPose2D(
				pdf->mean.x() + (*partsIt)[0], (pdf->mean.y() + (*partsIt)[1]),
				(pdf->mean.phi() + (*partsIt)[2]));
			itDest->d.normalizePhi();
		}
	}

	MRPT_END
}

void CPosePDFParticles::clear() { clearParticles(); }
void CPosePDFParticles::getMean(CPose2D& est_) const
{
	// Calc average on SE(2)
	const size_t n = m_particles.size();
	if (n)
	{
		mrpt::poses::SE_average<2> se_averager;
		for (size_t i = 0; i < n; i++)
		{
			double w = exp(m_particles[i].log_w);
			se_averager.append(m_particles[i].d, w);
		}
		se_averager.get_average(est_);
	}
	else
	{
		est_ = CPose2D();
	}
}

std::tuple<CMatrixDouble33, CPose2D> CPosePDFParticles::getCovarianceAndMean()
	const
{
	CMatrixDouble33 cov;
	CPose2D mean;
	cov.setZero();
	getMean(mean);

	size_t i, n = m_particles.size();
	double var_x = 0, var_y = 0, var_p = 0, var_xy = 0, var_xp = 0, var_yp = 0;
	double mean_phi = mean.phi();

	if (mean_phi < 0) mean_phi = M_2PI + mean_phi;

	double lin_w_sum = 0;

	for (i = 0; i < n; i++)
		lin_w_sum += exp(m_particles[i].log_w);
	if (lin_w_sum == 0) lin_w_sum = 1;

	for (i = 0; i < n; i++)
	{
		double w = exp(m_particles[i].log_w) / lin_w_sum;

		// Manage 1 PI range:
		double err_x = m_particles[i].d.x - mean.x();
		double err_y = m_particles[i].d.y - mean.y();
		double err_phi = math::wrapToPi(fabs(m_particles[i].d.phi - mean_phi));

		var_x += square(err_x) * w;
		var_y += square(err_y) * w;
		var_p += square(err_phi) * w;
		var_xy += err_x * err_y * w;
		var_xp += err_x * err_phi * w;
		var_yp += err_y * err_phi * w;
	}

	if (n < 2)
	{
		// Not enough information to estimate the variance:
	}
	else
	{
		// Unbiased estimation of variance:
		cov(0, 0) = var_x;
		cov(1, 1) = var_y;
		cov(2, 2) = var_p;

		cov(1, 0) = cov(0, 1) = var_xy;
		cov(2, 0) = cov(0, 2) = var_xp;
		cov(1, 2) = cov(2, 1) = var_yp;
	}
	return {cov, mean};
}

uint8_t CPosePDFParticles::serializeGetVersion() const { return 1; }
void CPosePDFParticles::serializeTo(mrpt::serialization::CArchive& out) const
{
	writeParticlesToStream(out);  // v1: changed CPose2D -> TPose2D
}
void CPosePDFParticles::serializeFrom(
	mrpt::serialization::CArchive& in, uint8_t version)
{
	switch (version)
	{
		case 0:
		{
			mrpt::bayes::CParticleFilterData<
				mrpt::poses::CPose2D, PARTICLE_STORAGE>
				old;
			old.readParticlesFromStream(in);
			m_particles.clear();
			std::transform(
				old.m_particles.begin(), old.m_particles.end(),
				std::back_inserter(m_particles),
				[](const auto& p) -> CParticleData {
					return CParticleData(p.d.asTPose(), p.log_w);
				});
		}
		break;
		case 1:
		{
			readParticlesFromStream(in);
		}
		break;
		default: MRPT_THROW_UNKNOWN_SERIALIZATION_VERSION(version);
	};
}

void CPosePDFParticles::resetDeterministic(
	const TPose2D& location, size_t particlesCount)
{
	if (particlesCount > 0) m_particles.resize(particlesCount);

	for (auto& p : m_particles)
	{
		p.d = location;
		p.log_w = .0;
	}
}

void CPosePDFParticles::resetUniform(
	const double x_min, const double x_max, const double y_min,
	const double y_max, const double phi_min, const double phi_max,
	const int particlesCount)
{
	MRPT_START
	if (particlesCount > 0) m_particles.resize(particlesCount);

	for (auto& p : m_particles)
	{
		p.d.x = getRandomGenerator().drawUniform(x_min, x_max);
		p.d.y = getRandomGenerator().drawUniform(y_min, y_max);
		p.d.phi = getRandomGenerator().drawUniform(phi_min, phi_max);
		p.log_w = 0;
	}
	MRPT_END
}

void CPosePDFParticles::resetAroundSetOfPoses(
	const std::vector<mrpt::math::TPose2D>& list_poses,
	const size_t num_particles_per_pose, const double spread_x,
	const double spread_y, const double spread_phi_rad)
{
	MRPT_START
	ASSERT_(!list_poses.empty());
	ASSERT_(num_particles_per_pose >= 1);

	const size_t N = list_poses.size() * num_particles_per_pose;

	clear();
	m_particles.resize(N);
	size_t i, nSpot;
	for (i = 0, nSpot = 0; nSpot < list_poses.size(); nSpot++)
	{
		const mrpt::math::TPose2D& p = list_poses[nSpot];
		for (size_t k = 0; k < num_particles_per_pose; k++, i++)
		{
			m_particles[i].d.x = getRandomGenerator().drawUniform(
				p.x - spread_x * 0.5, p.x + spread_x * 0.5);
			m_particles[i].d.y = getRandomGenerator().drawUniform(
				p.y - spread_y * 0.5, p.y + spread_y * 0.5);
			m_particles[i].d.phi = getRandomGenerator().drawUniform(
				p.phi - spread_phi_rad * 0.5, p.phi + spread_phi_rad * 0.5);
			m_particles[i].log_w = 0;
		}
	}
	ASSERT_EQUAL_(i, N);
	MRPT_END
}

bool CPosePDFParticles::saveToTextFile(const std::string& file) const
{
	using namespace mrpt::system;

	FILE* f = os::fopen(file.c_str(), "wt");
	if (!f) return false;

	os::fprintf(f, "%% x  y  yaw[rad] log_weight\n");

	for (const auto& p : m_particles)
		os::fprintf(f, "%f %f %f %e\n", p.d.x, p.d.y, p.d.phi, p.log_w);

	os::fclose(f);
	return true;
}

TPose2D CPosePDFParticles::getParticlePose(size_t i) const
{
	return m_particles[i].d;
}

void CPosePDFParticles::changeCoordinatesReference(
	const CPose3D& newReferenceBase_)
{
	const TPose2D newReferenceBase = CPose2D(newReferenceBase_).asTPose();
	for (auto& p : m_particles)
		p.d = newReferenceBase + p.d;
}

void CPosePDFParticles::drawSingleSample(CPose2D& outPart) const
{
	const double uni = getRandomGenerator().drawUniform(0.0, 0.9999);
	double cum = 0;

	for (auto& p : m_particles)
	{
		cum += exp(p.log_w);
		if (uni <= cum)
		{
			outPart = CPose2D(p.d);
			return;
		}
	}

	// Might not come here normally:
	outPart = CPose2D(m_particles.rbegin()->d);
}

void CPosePDFParticles::operator+=(const TPose2D& Ap)
{
	for (auto& p : m_particles)
		p.d = p.d + Ap;
}

void CPosePDFParticles::append(CPosePDFParticles& o)
{
	for (auto& p : o.m_particles)
		m_particles.emplace_back(p);
	normalizeWeights();
}

void CPosePDFParticles::inverse(CPosePDF& o) const
{
	MRPT_START
	ASSERT_(o.GetRuntimeClass() == CLASS_ID(CPosePDFParticles));
	auto* out = dynamic_cast<CPosePDFParticles*>(&o);

	out->copyFrom(*this);
	TPose2D nullPose(0, 0, 0);

	for (auto& p : out->m_particles)
		p.d = nullPose - p.d;

	MRPT_END
}

mrpt::math::TPose2D CPosePDFParticles::getMostLikelyParticle() const
{
	mrpt::math::TPose2D ret{0, 0, 0};
	double max_w = -std::numeric_limits<double>::max();
	for (const auto& p : m_particles)
	{
		if (p.log_w > max_w)
		{
			ret = p.d;
			max_w = p.log_w;
		}
	}
	return ret;
}

void CPosePDFParticles::bayesianFusion(
	[[maybe_unused]] const CPosePDF& p1, [[maybe_unused]] const CPosePDF& p2,
	[[maybe_unused]] const double minMahalanobisDistToDrop)
{
	THROW_EXCEPTION("Not implemented yet!");
}

double CPosePDFParticles::evaluatePDF_parzen(
	const double x, const double y, const double phi, const double stdXY,
	const double stdPhi) const
{
	double ret = 0;
	for (const auto& p : m_particles)
	{
		double difPhi = math::wrapToPi(phi - p.d.phi);
		ret += exp(p.log_w) *
			math::normalPDF(
				   std::sqrt(square(p.d.x - x) + square(p.d.y - y)), 0, stdXY) *
			math::normalPDF(std::abs(difPhi), 0, stdPhi);
	}
	return ret;
}

void CPosePDFParticles::saveParzenPDFToTextFile(
	const char* fileName, const double x_min, const double x_max,
	const double y_min, const double y_max, const double phi,
	const double stepSizeXY, const double stdXY, const double stdPhi) const
{
	std::string buf;

	for (double y = y_min; y < y_max; y += stepSizeXY)
		for (double x = x_min; x < x_max; x += stepSizeXY)
			buf += mrpt::format(
				"%f ", evaluatePDF_parzen(x, y, phi, stdXY, stdPhi));
	buf += "\n";

	std::ofstream f(fileName);
	if (!f.is_open()) return;
	f << buf;
}

std::string CPosePDFParticles::asString() const
{
	std::stringstream ss;
	const auto [c, m] = getCovarianceAndMean();
	ss << "mrpt::poses::CPosePDFParticles object with " << size()
	   << " particles, mean=" << m.asString() << " cov=" << c.inMatlabFormat()
	   << " ESS=" << ESS();
	return ss.str();
}
