/* +------------------------------------------------------------------------+
   |                     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/TPose3D.h>
#include <mrpt/poses/CPose3D.h>
#include <mrpt/poses/CPose3DPDF.h>
#include <mrpt/poses/CPose3DPDFGaussian.h>
#include <mrpt/poses/CPose3DPDFGaussianInf.h>
#include <mrpt/poses/CPose3DPDFParticles.h>
#include <mrpt/poses/CPose3DPDFSOG.h>
#include <mrpt/poses/CPose3DQuatPDF.h>
#include <mrpt/poses/CPosePDFGaussian.h>
#include <mrpt/poses/CPosePDFGaussianInf.h>
#include <mrpt/poses/CPosePDFParticles.h>
#include <mrpt/poses/CPosePDFSOG.h>
#include <mrpt/serialization/CArchive.h>

#include <Eigen/Dense>

using namespace mrpt::poses;
using namespace mrpt::math;
using namespace std;

IMPLEMENTS_VIRTUAL_SERIALIZABLE(CPose3DPDF, CSerializable, mrpt::poses)

// copy covariances: x, y
// Move covariances: phi -> yaw
template <class MAT33, class MAT66>
void cov2to3(const MAT33& c2d, MAT66& c3d)
{
	c3d.setZero();
	c3d.asEigen().template block<2, 2>(0, 0) =
		c2d.asEigen().template block<2, 2>(0, 0);
	c3d(3, 3) = c2d(2, 2);
	c3d(0, 3) = c3d(3, 0) = c2d(0, 2);
	c3d(1, 3) = c3d(3, 1) = c2d(1, 2);
}

/*---------------------------------------------------------------
					copyFrom2D
  ---------------------------------------------------------------*/
CPose3DPDF* CPose3DPDF::createFrom2D(const CPosePDF& o)
{
	MRPT_START

	if (o.GetRuntimeClass() == CLASS_ID(CPosePDFGaussian))
	{
		auto* newObj = new CPose3DPDFGaussian();
		const auto* obj = dynamic_cast<const CPosePDFGaussian*>(&o);
		ASSERT_(obj != nullptr);
		newObj->mean = CPose3D(obj->mean);
		cov2to3(obj->cov, newObj->cov);
		return newObj;
	}
	else if (o.GetRuntimeClass() == CLASS_ID(CPosePDFGaussianInf))
	{
		auto* newObj = new CPose3DPDFGaussianInf();
		const auto* obj = dynamic_cast<const CPosePDFGaussianInf*>(&o);
		ASSERT_(obj != nullptr);

		newObj->mean = CPose3D(obj->mean);
		cov2to3(obj->cov_inv, newObj->cov_inv);
		return newObj;
	}
	else if (o.GetRuntimeClass() == CLASS_ID(CPosePDFParticles))
	{
		const auto* obj = dynamic_cast<const CPosePDFParticles*>(&o);
		ASSERT_(obj != nullptr);
		auto* newObj = new CPose3DPDFParticles(obj->size());

		CPosePDFParticles::CParticleList::const_iterator it1;
		CPose3DPDFParticles::CParticleList::iterator it2;
		for (it1 = obj->m_particles.begin(), it2 = newObj->m_particles.begin();
			 it1 != obj->m_particles.end(); ++it1, ++it2)
		{
			it2->log_w = it1->log_w;
			it2->d = TPose3D(it1->d);
		}

		return newObj;
	}
	else if (o.GetRuntimeClass() == CLASS_ID(CPosePDFSOG))
	{
		const auto* obj = dynamic_cast<const CPosePDFSOG*>(&o);
		ASSERT_(obj != nullptr);
		auto* newObj = new CPose3DPDFSOG(obj->size());

		CPosePDFSOG::const_iterator it1;
		CPose3DPDFSOG::iterator it2;

		for (it1 = obj->begin(), it2 = newObj->begin(); it1 != obj->end();
			 ++it1, ++it2)
		{
			it2->log_w = it1->log_w;
			it2->val.mean.setFromValues(
				it1->mean.x(), it1->mean.y(), 0, it1->mean.phi(), 0, 0);

			it2->val.cov.setZero();

			it2->val.cov(0, 0) = it1->cov(0, 0);
			it2->val.cov(1, 1) = it1->cov(1, 1);
			it2->val.cov(3, 3) = it1->cov(2, 2);  // yaw <- phi

			it2->val.cov(0, 1) = it2->val.cov(1, 0) = it1->cov(0, 1);

			it2->val.cov(0, 3) = it2->val.cov(3, 0) = it1->cov(0, 2);

			it2->val.cov(1, 3) = it2->val.cov(3, 1) = it1->cov(1, 2);
		}

		return newObj;
	}
	else
		THROW_EXCEPTION("Class of object not supported by this method!");

	MRPT_END
}

/*---------------------------------------------------------------
					jacobiansPoseComposition
 ---------------------------------------------------------------*/
void CPose3DPDF::jacobiansPoseComposition(
	const CPose3D& x, const CPose3D& u, CMatrixDouble66& df_dx,
	CMatrixDouble66& df_du)
{
	// See this techical report:
	// https://www.mrpt.org/6D_poses:equivalences_compositions_and_uncertainty

	// Direct equations (for the covariances) in yaw-pitch-roll are too complex.
	//  Make a way around them and consider instead this path:
	//
	//      X(6D)       U(6D)
	//        |           |
	//        v           v
	//      X(7D)       U(7D)
	//        |           |
	//        +--- (+) ---+
	//              |
	//              v
	//            RES(7D)
	//              |
	//              v
	//            RES(6D)
	//

	// FUNCTION = f_quat2eul(  f_quat_comp(  f_eul2quat(x) , f_eul2quat(u)  ) )
	//  Jacobians chain rule:
	//
	//  JACOB_dx = J_Q2E (6x7) * quat_df_dx (7x7) * J_E2Q_dx (7x6)
	//  JACOB_du = J_Q2E (6x7) * quat_df_du (7x7) * J_E2Q_du (7x6)

	// J_E2Q_dx:
	CMatrixFixed<double, 7, 6> J_E2Q_dx;  // Init to zeros
	{
		CMatrixFixed<double, 4, 3> dq_dr_sub(UNINITIALIZED_MATRIX);
		CQuaternionDouble q_dumm(UNINITIALIZED_QUATERNION);
		x.getAsQuaternion(q_dumm, dq_dr_sub);
		J_E2Q_dx(0, 0) = J_E2Q_dx(1, 1) = J_E2Q_dx(2, 2) = 1;
		J_E2Q_dx.insertMatrix(3, 3, dq_dr_sub);
	}

	// J_E2Q_du:
	CMatrixFixed<double, 7, 6> J_E2Q_du;  // Init to zeros
	{
		CMatrixFixed<double, 4, 3> dq_dr_sub(UNINITIALIZED_MATRIX);
		CQuaternionDouble q_dumm(UNINITIALIZED_QUATERNION);
		u.getAsQuaternion(q_dumm, dq_dr_sub);
		J_E2Q_du(0, 0) = J_E2Q_du(1, 1) = J_E2Q_du(2, 2) = 1;
		J_E2Q_du.insertMatrix(3, 3, dq_dr_sub);
	}

	// quat_df_dx & quat_df_du
	CMatrixDouble77 quat_df_dx(UNINITIALIZED_MATRIX),
		quat_df_du(UNINITIALIZED_MATRIX);
	const CPose3DQuat quat_x(x);
	const CPose3DQuat quat_u(u);

	CPose3DQuatPDF::jacobiansPoseComposition(
		quat_x,	 // x
		quat_u,	 // u
		quat_df_dx, quat_df_du);

	// And now J_Q2E:
	//         [  I_3   |    0         ]
	// J_Q2E = [ -------+------------- ]
	//         [  0     | dr_dq_angles ]
	//
	CMatrixFixed<double, 6, 7> J_Q2E;  // Init to zeros
	J_Q2E(0, 0) = J_Q2E(1, 1) = J_Q2E(2, 2) = 1;
	{
		// The end result of the pose composition, as a quaternion:
		CQuaternionDouble q_xu(UNINITIALIZED_QUATERNION);
		q_xu.crossProduct(quat_x.quat(), quat_u.quat());

		// Compute the jacobian:
		CMatrixFixed<double, 3, 4> dr_dq_sub_aux(UNINITIALIZED_MATRIX);
		double yaw, pitch, roll;
		q_xu.rpy_and_jacobian(roll, pitch, yaw, &dr_dq_sub_aux, false);

		CMatrixDouble44 dnorm_dq(UNINITIALIZED_MATRIX);
		q_xu.normalizationJacobian(dnorm_dq);

		CMatrixFixed<double, 3, 4> dr_dq_sub(UNINITIALIZED_MATRIX);
		dr_dq_sub.asEigen() = dr_dq_sub_aux * dnorm_dq;

		J_Q2E.insertMatrix(3, 3, dr_dq_sub);
	}

	// And finally:
	//  JACOB_dx = J_Q2E (6x7) * quat_df_dx (7x7) * J_E2Q_dx (7x6)
	//  JACOB_du = J_Q2E (6x7) * quat_df_du (7x7) * J_E2Q_du (7x6)
	df_dx.asEigen() =
		J_Q2E.asEigen() * quat_df_dx.asEigen() * J_E2Q_dx.asEigen();
	df_du.asEigen() =
		J_Q2E.asEigen() * quat_df_du.asEigen() * J_E2Q_du.asEigen();
}
