/* +------------------------------------------------------------------------+
   |                     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 <mrpt/graphs.h>
#include <mrpt/graphslam/levmarq.h>
#include <mrpt/graphslam/types.h>
#include <mrpt/io/CMemoryStream.h>
#include <mrpt/random.h>

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

template <class my_graph_t>
class GraphSlamLevMarqTest
{
   public:
	// adds a new edge to the graph. The edge is annotated with the relative
	// position of the two nodes
	static void addEdge(
		TNodeID from, TNodeID to,
		const typename my_graph_t::global_poses_t& real_poses,
		my_graph_t& graph)
	{
		auto RelativePose =
			real_poses.find(to)->second - real_poses.find(from)->second;

		// Add information matrix if present:
		if constexpr (my_graph_t::edge_t::is_PDF())
		{
			const auto N = my_graph_t::edge_t::state_length;
			using InfMat = mrpt::math::CMatrixFixed<double, N, N>;
			auto mat = getRandomGenerator()
						   .drawDefinitePositiveMatrix<
							   InfMat, mrpt::math::CVectorFixed<double, N>>(
							   N, 2.0 /*std*/, 1.0 /*diagonal offset*/);
			graph.insertEdge(
				from, to, typename my_graph_t::edge_t(RelativePose, mat));
		}
		else
		{
			graph.insertEdge(from, to, RelativePose);
		}
	}

	// The graph: nodes + edges:
	static void create_ring_path(
		my_graph_t& graph, size_t N_VERTEX = 50, double DIST_THRES = 7,
		double NODES_XY_MAX = 20)
	{
		// The global poses of each node (without covariance):
		typename my_graph_t::global_poses_t real_node_poses;

		// ----------------------------
		// Create a random graph:
		// ----------------------------
		// Level of noise in nodes initial positions:
		const double STD_NOISE_NODE_XYZ = 0.5;
		const double STD_NOISE_NODE_ANG = 5.0_deg;

		// Level of noise in edges:
		const double STD_NOISE_EDGE_XYZ = 0;  // 0.01;
		const double STD_NOISE_EDGE_ANG = 0;  // 0.1_deg;

		for (TNodeID j = 0; j < N_VERTEX; j++)
		{
			static double ang = 2 * M_PI / N_VERTEX;
			const double R = NODES_XY_MAX + 2 * (j % 2 ? 1 : -1);
			CPose2D p(R * cos(ang * j), R * sin(ang * j), ang);

			// Save real pose:
			real_node_poses[j] = p;

			// Copy the nodes to the graph, and add some noise:
			graph.nodes[j] = p;
		}

		// Add some edges
		for (TNodeID i = 0; i < N_VERTEX; i++)
		{
			for (TNodeID j = i + 1; j < N_VERTEX; j++)
			{
				if (real_node_poses[i].distanceTo(real_node_poses[j]) <
					DIST_THRES)
					addEdge(i, j, real_node_poses, graph);
			}
		}

		// Cross-links:
		addEdge(0, N_VERTEX / 2, real_node_poses, graph);

		// The root node (the origin of coordinates):
		graph.root = TNodeID(0);

		// This is the ground truth graph (make a copy for later use):
		const my_graph_t graph_GT = graph;

		// Add noise to edges & nodes:
		for (auto& edge : graph.edges)
		{
			// edge_t::type_value: for edge_t a plain pose, type_value is the
			// edge_t itself. For poses+uncertainty, it is the underlying pose
			// type.
			if constexpr (my_graph_t::edge_t::type_value::size() == 6)
			{
				const auto noise = CPose3D(
					getRandomGenerator().drawGaussian1D(0, STD_NOISE_EDGE_XYZ),
					getRandomGenerator().drawGaussian1D(0, STD_NOISE_EDGE_XYZ),
					getRandomGenerator().drawGaussian1D(0, STD_NOISE_EDGE_XYZ),
					getRandomGenerator().drawGaussian1D(0, STD_NOISE_EDGE_ANG),
					getRandomGenerator().drawGaussian1D(0, STD_NOISE_EDGE_ANG),
					getRandomGenerator().drawGaussian1D(0, STD_NOISE_EDGE_ANG));

				if constexpr (my_graph_t::edge_t::is_PDF())
				{
					const auto N = my_graph_t::edge_t::state_length;
					using InfMat = mrpt::math::CMatrixFixed<double, N, N>;
					InfMat mat;
					mat.setIdentity();
					edge.second += typename my_graph_t::edge_t(noise, mat);
				}
				else
				{
					edge.second += typename my_graph_t::edge_t(noise);
				}
			}
			else
			{
				const auto noise = CPose2D(
					getRandomGenerator().drawGaussian1D(0, STD_NOISE_EDGE_XYZ),
					getRandomGenerator().drawGaussian1D(0, STD_NOISE_EDGE_XYZ),
					getRandomGenerator().drawGaussian1D(0, STD_NOISE_EDGE_ANG));
				if constexpr (my_graph_t::edge_t::is_PDF())
				{
					const auto N = my_graph_t::edge_t::state_length;
					using InfMat = mrpt::math::CMatrixFixed<double, N, N>;
					InfMat mat;
					mat.setIdentity();
					edge.second += typename my_graph_t::edge_t(noise, mat);
				}
				else
				{
					edge.second += typename my_graph_t::edge_t(noise);
				}
			}
		}

		for (auto itNode = graph.nodes.begin(); itNode != graph.nodes.end();
			 ++itNode)
			if (itNode->first != graph.root)
				itNode
					->second += typename my_graph_t::edge_t::type_value(CPose3D(
					getRandomGenerator().drawGaussian1D(0, STD_NOISE_NODE_XYZ),
					getRandomGenerator().drawGaussian1D(0, STD_NOISE_NODE_XYZ),
					getRandomGenerator().drawGaussian1D(0, STD_NOISE_NODE_XYZ),
					getRandomGenerator().drawGaussian1D(0, STD_NOISE_NODE_ANG),
					getRandomGenerator().drawGaussian1D(0, STD_NOISE_NODE_ANG),
					getRandomGenerator().drawGaussian1D(
						0, STD_NOISE_NODE_ANG)));
	}
};
