/* +------------------------------------------------------------------------+
   |                     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 <gtest/gtest.h>
#include <mrpt/math/CMatrixFixed.h>
#include <mrpt/math/num_jacobian.h>
#include <mrpt/poses/CPose2D.h>
#include <mrpt/poses/CPose3D.h>
#include <mrpt/poses/Lie/SE.h>

#include <Eigen/Dense>

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

// List of "random" poses to test with (x,y,z,yaw,pitch,roll) (angles in
// degrees)
static const std::vector<mrpt::poses::CPose3D> ptc = {
	{.0, .0, .0, .0_deg, .0_deg, .0_deg},
	{1.0, 2.0, 3.0, .0_deg, .0_deg, .0_deg},
	{1.0, 2.0, 3.0, 10.0_deg, .0_deg, .0_deg},
	{1.0, 2.0, 3.0, .0_deg, 1.0_deg, .0_deg},
	{1.0, 2.0, 3.0, .0_deg, .0_deg, 1.0_deg},
	{-6.0, 2.0, 3.0, 40.0_deg, 20.0_deg, 15.0_deg},
	{1.0, 8.0, 0.0, -45.0_deg, 10.0_deg, 70.0_deg},
	{1.0, 2.0, 3.0, 89.0_deg, .5_deg, 1.0_deg},
	{1.0, -2.0, 0.4, -89.0_deg, .5_deg, 1.0_deg},
	{1.0, 2.0, 3.0, .1_deg, 89.0_deg, 1.0_deg},
	{1.0, -2.0, 0.4, .1_deg, -88.0_deg, 1.0_deg},
	{1.0, 2.0, 3.0, .1_deg, .9_deg, 178.0_deg},
	{1.0, -2.0, 0.4, .1_deg, .9_deg, -179.1_deg},
};

template <class POSE_TYPE>
class SE_traits_tests : public ::testing::Test
{
   protected:
	void SetUp() override {}
	void TearDown() override {}
	using SE_TYPE = mrpt::poses::Lie::SE<POSE_TYPE::rotation_dimensions>;

	struct TParams
	{
		typename SE_TYPE::type P1, D, P2;
	};

	template <unsigned int MAT_LEN>
	struct TParamsMat
	{
		CVectorFixedDouble<MAT_LEN> Avec, Bvec;
	};

	static void func_numeric_DinvP1InvP2(
		const CVectorFixedDouble<2 * SE_TYPE::DOFs>& x, const TParams& params,
		CVectorFixedDouble<SE_TYPE::DOFs>& Y)
	{
		typename SE_TYPE::tangent_vector eps1, eps2;
		for (size_t i = 0; i < SE_TYPE::DOFs; i++)
		{
			eps1[i] = x[0 + i];
			eps2[i] = x[SE_TYPE::DOFs + i];
		}

		const POSE_TYPE incr1 = SE_TYPE::exp(eps1);
		const POSE_TYPE incr2 = SE_TYPE::exp(eps2);

		const POSE_TYPE P1 = params.P1 + incr1;
		const POSE_TYPE P2 = params.P2 + incr2;
		const POSE_TYPE& Pd = params.D;

		CMatrixDouble44 P1_inv_hm, P2hm, Pd_inv_hm;
		P1.getInverseHomogeneousMatrix(P1_inv_hm);
		P2.getHomogeneousMatrix(P2hm);
		Pd.getInverseHomogeneousMatrix(Pd_inv_hm);

		const CPose3D DinvP1invP2_(CMatrixDouble44(
			Pd_inv_hm.asEigen() * P1_inv_hm.asEigen() * P2hm.asEigen()));
		const POSE_TYPE DinvP1invP2(DinvP1invP2_);

		// Pseudo-logarithm:
		Y = SE_TYPE::log(DinvP1invP2);
	}

	void test_jacobs_DinvP1InvP2(
		const CPose3D& P1_, const CPose3D& Pd_, const CPose3D& P2_)
	{
		const POSE_TYPE P1(P1_);
		const POSE_TYPE Pd(Pd_);
		const POSE_TYPE P2(P2_);
		const POSE_TYPE Pdinv = -Pd;

		static const int DIMS = SE_TYPE::DOFs;

		// Theoretical results:
		CMatrixFixed<double, DIMS, DIMS> J1, J2;
		SE_TYPE::jacob_dDinvP1invP2_de1e2(Pdinv, P1, P2, J1, J2);

		// Numerical approx:
		CMatrixFixed<double, DIMS, DIMS> num_J1, num_J2;
		{
			CVectorFixedDouble<2 * DIMS> x_mean;
			for (int i = 0; i < DIMS + DIMS; i++)
				x_mean[i] = 0;

			TParams params;
			params.P1 = P1;
			params.D = Pd;
			params.P2 = P2;

			CVectorFixedDouble<DIMS + DIMS> x_incrs;
			x_incrs.fill(1e-4);
			CMatrixDouble numJacobs;
			mrpt::math::estimateJacobian(
				x_mean,
				std::function<void(
					const CVectorFixedDouble<2 * SE_TYPE::DOFs>& x,
					const TParams& params,
					CVectorFixedDouble<SE_TYPE::DOFs>& Y)>(
					&func_numeric_DinvP1InvP2),
				x_incrs, params, numJacobs);

			num_J1 = numJacobs.asEigen().block<DIMS, DIMS>(0, 0);
			num_J2 = numJacobs.asEigen().block<DIMS, DIMS>(0, DIMS);
		}

		const double max_eror = 1e-3;

		EXPECT_NEAR(0, (num_J1 - J1).sum_abs(), max_eror)
			<< std::setprecision(3) << "p1: " << P1 << endl
			<< "d: " << Pd << endl
			<< "p2: " << P2 << endl
			<< "Numeric J1:\n"
			<< num_J1 << endl
			<< "Implemented J1:\n"
			<< J1 << endl
			<< "Error:\n"
			<< J1 - num_J1 << endl;

		EXPECT_NEAR(0, (num_J2 - J2).sum_abs(), max_eror)
			<< "p1: " << P1 << endl
			<< "d: " << Pd << endl
			<< "p2: " << P2 << endl
			<< "Numeric J2:\n"
			<< num_J2 << endl
			<< "Implemented J2:\n"
			<< J2 << endl
			<< "Error:\n"
			<< J2 - num_J2 << endl;
	}

	static void func_numeric_dAB_dA(
		const CVectorFixedDouble<SE_TYPE::MANIFOLD_DIM>& x,
		const TParamsMat<SE_TYPE::MANIFOLD_DIM>& params,
		CVectorFixedDouble<SE_TYPE::MANIFOLD_DIM>& Y)
	{
		const POSE_TYPE A = SE_TYPE::fromManifoldVector(params.Avec + x);
		const POSE_TYPE B = SE_TYPE::fromManifoldVector(params.Bvec);
		Y = SE_TYPE::asManifoldVector(A + B);
	}

	static void func_numeric_dAB_dB(
		const CVectorFixedDouble<SE_TYPE::MANIFOLD_DIM>& x,
		const TParamsMat<SE_TYPE::MANIFOLD_DIM>& params,
		CVectorFixedDouble<SE_TYPE::MANIFOLD_DIM>& Y)
	{
		const POSE_TYPE A = SE_TYPE::fromManifoldVector(params.Avec);
		const POSE_TYPE B = SE_TYPE::fromManifoldVector(params.Bvec + x);
		Y = SE_TYPE::asManifoldVector(A + B);
	}

	void test_jacobs_AB(const CPose3D& A_, const CPose3D& B_)
	{
		const POSE_TYPE A(A_);
		const POSE_TYPE B(B_);

		constexpr auto N = SE_TYPE::MANIFOLD_DIM;

		// Theoretical results:
		const auto dAB_A = SE_TYPE::jacob_dAB_dA(A, B);
		const auto dAB_B = SE_TYPE::jacob_dAB_dB(A, B);

		// Numerical approx: dAB_A
		CMatrixFixed<double, N, N> num_dAB_A;
		{
			CVectorFixedDouble<N> x_mean;
			x_mean.fill(0);

			TParamsMat<N> params;
			params.Avec = SE_TYPE::asManifoldVector(A);
			params.Bvec = SE_TYPE::asManifoldVector(B);

			CVectorFixedDouble<N> x_incrs;
			x_incrs.fill(1e-4);
			CMatrixDouble numJacobs;
			mrpt::math::estimateJacobian(
				x_mean,
				std::function<void(
					const CVectorFixedDouble<N>& x, const TParamsMat<N>& params,
					CVectorFixedDouble<N>& Y)>(&func_numeric_dAB_dA),
				x_incrs, params, numJacobs);

			num_dAB_A = numJacobs;
		}

		// Numerical approx: dAB_B
		CMatrixFixed<double, N, N> num_dAB_B;
		{
			CVectorFixedDouble<N> x_mean;
			x_mean.fill(0);

			TParamsMat<N> params;
			params.Avec = SE_TYPE::asManifoldVector(A);
			params.Bvec = SE_TYPE::asManifoldVector(B);

			CVectorFixedDouble<N> x_incrs;
			x_incrs.fill(1e-4);
			CMatrixDouble numJacobs;
			mrpt::math::estimateJacobian(
				x_mean,
				std::function<void(
					const CVectorFixedDouble<N>& x, const TParamsMat<N>& params,
					CVectorFixedDouble<N>& Y)>(&func_numeric_dAB_dB),
				x_incrs, params, numJacobs);

			num_dAB_B = numJacobs;
		}

		const double max_eror = 1e-3;

		EXPECT_NEAR(0, (num_dAB_A - dAB_A).sum_abs(), max_eror)
			<< std::setprecision(3) << "A: " << A << endl
			<< "B: " << B << endl
			<< "Numeric dAB_A:\n"
			<< num_dAB_A << endl
			<< "Implemented dAB_A:\n"
			<< dAB_A << endl
			<< "Error:\n"
			<< dAB_A - num_dAB_A << endl;

		EXPECT_NEAR(0, (num_dAB_B - dAB_B).sum_abs(), max_eror)
			<< std::setprecision(3) << "A: " << A << endl
			<< "B: " << B << endl
			<< "Numeric dAB_B:\n"
			<< num_dAB_B << endl
			<< "Implemented dAB_B:\n"
			<< dAB_B << endl
			<< "Error:\n"
			<< dAB_B - num_dAB_B << endl;
	}

	void tests_jacobs_DinvP1InvP2()
	{
		for (const auto& p1 : ptc)
			for (const auto& p2 : ptc)
				for (const auto& pd : ptc)
					test_jacobs_DinvP1InvP2(p1, pd, p2);
	}

	void tests_jacobs_dAB_dAB()
	{
		for (const auto& p1 : ptc)
			for (const auto& p2 : ptc)
				test_jacobs_AB(p1, p2);
	}
};

using SE3_traits_tests = SE_traits_tests<mrpt::poses::CPose3D>;
using SE2_traits_tests = SE_traits_tests<mrpt::poses::CPose2D>;

TEST_F(SE3_traits_tests, SE3_jacobs_DinvP1InvP2) { tests_jacobs_DinvP1InvP2(); }
TEST_F(SE3_traits_tests, SE3_jacobs_dAB_dAB) { tests_jacobs_dAB_dAB(); }

TEST_F(SE2_traits_tests, SE2_jacobs_DinvP1InvP2) { tests_jacobs_DinvP1InvP2(); }
TEST_F(SE2_traits_tests, SE2_jacobs_dAB_dAB) { tests_jacobs_dAB_dAB(); }
