/* +------------------------------------------------------------------------+
   |                     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          |
   +------------------------------------------------------------------------+ */

/*---------------------------------------------------------------
	APPLICATION: Graph-SLAM hub application
	PURPOSE: Visualization of .graph files and evaluation of
			 Graph-SLAM methods.

	See the "--help" output for list of supported operations and
	 further instructions, the man-page or online docs at:
		 https://www.mrpt.org/list-of-mrpt-apps/application-graph-slam
 ---------------------------------------------------------------*/
#include <mrpt/3rdparty/tclap/CmdLine.h>
#include <mrpt/graphs.h>
#include <mrpt/graphslam/levmarq.h>
#include <mrpt/system/datetime.h>

using namespace mrpt;
using namespace mrpt::graphslam;
using namespace mrpt::graphs;
using namespace mrpt::poses;
using namespace mrpt::opengl;
using namespace mrpt::system;
using namespace mrpt::math;
using namespace std;

// Declarations:
#define VERBOSE_COUT                                                           \
	if (verbose) std::cout << "[graph-slam] "

#define DECLARE_OP_FUNCTION(_NAME)                                             \
	void _NAME(                                                                \
		const std::string& in_file, bool is3D, TCLAP::CmdLine& cmdline,        \
		bool verbose)

#define IMPLEMENT_OP_FUNCTION(_NAME)                                           \
	template <class GRAPHTYPE>                                                 \
	void _NAME##impl(                                                          \
		const std::string& in_file, bool is3D, TCLAP::CmdLine& cmdline,        \
		bool verbose);                                                         \
	void _NAME(                                                                \
		const std::string& in_file, bool is3D, TCLAP::CmdLine& cmdline,        \
		bool verbose)                                                          \
	{                                                                          \
		if (!is3D)                                                             \
			_NAME##impl<CNetworkOfPoses2D>(in_file, is3D, cmdline, verbose);   \
		else                                                                   \
			_NAME##impl<CNetworkOfPoses3D>(in_file, is3D, cmdline, verbose);   \
	}                                                                          \
	template <class GRAPHTYPE>                                                 \
	void _NAME##impl(                                                          \
		const std::string& in_file, bool is3D, TCLAP::CmdLine& cmdline,        \
		bool verbose)

/**
 * http://stackoverflow.com/questions/3982470/what-does-typedef-void-something-mean
 * Define a pointer-to-function type "TOperationFunctor". TOperationFunctor
 * returns void and takes 4 input parameters as explained in the latter
 * parentheses
 */
using TOperationFunctor = void (*)(
	const std::string& in_file, bool is3D, TCLAP::CmdLine& cmdline,
	bool verbose);

/** Forward declaration.
 * display_graph function template provided in display_graph.cpp
 */
template <class GRAPHTYPE>
void display_graph(const GRAPHTYPE& g);

DECLARE_OP_FUNCTION(op_view);
DECLARE_OP_FUNCTION(op_dijkstra);
DECLARE_OP_FUNCTION(op_levmarq);
DECLARE_OP_FUNCTION(op_info);

// Declare the supported command line switches ===========
TCLAP::CmdLine cmd("graph-slam", ' ', MRPT_getVersion().c_str());

TCLAP::ValueArg<std::string> arg_input_file(
	"i", "input", "Input file (required) (*.graph,*.graphbin)", true, "",
	"test.graph", cmd);
TCLAP::ValueArg<std::string> arg_output_file(
	"o", "output", "Output file (optional) (*.graph,*.graphbin)", false, "",
	"result.graph", cmd);

TCLAP::SwitchArg arg_2d(
	"", "2d", "Use 2D poses (Must use exactly one of --2d and --3d)", cmd,
	false);
TCLAP::SwitchArg arg_3d(
	"", "3d", "Use 3D poses (Must use exactly one of --2d and --3d)", cmd,
	false);

TCLAP::SwitchArg arg_view(
	"", "view",
	"Op: Visualize the graph in a 3D view. If used alone, represent VERTEX2 or "
	"VERTEX3 poses directly as stored in the input file. "
	"If used together with another operation, the final obtained graph after "
	"the operation will be shown, not the input original one.\n",
	cmd, false);

TCLAP::SwitchArg arg_quiet("q", "quiet", "Terse output", cmd, false);
TCLAP::ValueArg<int> arg_max_iters(
	"", "max-iters", "Maximum number of iterations (optional)", false, 100, "N",
	cmd);
TCLAP::ValueArg<double> arg_initial_lambda(
	"", "initial-lambda", "Initial lambda parameter (optional, lev-marq)",
	false, 0, "val", cmd);
TCLAP::SwitchArg arg_no_span(
	"", "no-span", "Don't use dijkstra initial spanning tree guess (optional)",
	cmd, false);

// ======================================================================
//     main() of graph-slam
// ======================================================================
int main(int argc, char** argv)
{
	vector<std::unique_ptr<TCLAP::Arg>> arg_ops;
	int ret_val = 0;

	try
	{
		// --------------- List of possible operations ---------------
		// Only one of the following operators have to be specified by the user
		map<string, TOperationFunctor> ops_functors;

		arg_ops.push_back(std::make_unique<TCLAP::SwitchArg>(
			"", "levmarq",
			"Op: Optimizes the graph with sparse Levenberg-Marquartd using "
			"global coordinates (via "
			"mrpt::graphslam::optimize_graph_spa_levmarq).\n"
			"   Can be used together with: --view, --output, --max-iters, "
			"--no-span, --initial-lambda",
			cmd, false));
		ops_functors["levmarq"] = &op_levmarq;

		arg_ops.push_back(std::make_unique<TCLAP::SwitchArg>(
			"", "dijkstra",
			"Op: Executes CNetworkOfPoses::dijkstra_nodes_estimate() to "
			"estimate the global pose of nodes from a Dijkstra tree and "
			"the edge relative poses.\n"
			"   Can be used together with: --view, --output",
			cmd, false));
		ops_functors["dijkstra"] = &op_dijkstra;

		arg_ops.push_back(std::make_unique<TCLAP::SwitchArg>(
			"", "info",
			"Op: Loads the graph and displays statistics and information "
			"on it.\n",
			cmd, false));
		ops_functors["info"] = &op_info;
		// --------------- End of list of possible operations --------

		// Parse arguments:
		if (!cmd.parse(argc, argv))
			throw std::runtime_error("");  // should exit.

		// Exactly 1 or --2d & --3d must be specified:
		if ((arg_2d.isSet() && arg_3d.isSet()) ||
			(!arg_2d.isSet() && !arg_3d.isSet()))
			throw std::runtime_error("Exactly one --2d or --3d must be used.");

		const bool is3d = arg_3d.isSet();

		string input_file = arg_input_file.getValue();
		const bool verbose = !arg_quiet.getValue();

		// Check the selected operation:
		//  Only one of the ops should be selected:
		string selected_op;
		for (auto& arg_op : arg_ops)
		{
			if (arg_op->isSet())
			{
				if (selected_op.empty()) { selected_op = arg_op->getName(); }
				else
					throw std::runtime_error(
						"Exactly one operation must be indicated on command "
						"line.\n"
						"Use --help to see the list of possible operations.");
			}
		}

		// The "--view" argument needs a bit special treatment:
		if (selected_op.empty())
		{
			if (!arg_view.isSet())
				throw std::runtime_error(
					"Don't know what to do: No operation was indicated.\n"
					"Use --help to see the list of possible operations.");
			else
			{
				VERBOSE_COUT << "Operation to perform: "
								"view"
							 << endl;
				op_view(input_file, is3d, cmd, verbose);
			}
		}
		else
		{
			VERBOSE_COUT << "Operation to perform: " << selected_op << endl;

			// ------------------------------------
			//  EXECUTE THE REQUESTED OPERATION
			// ------------------------------------
			ASSERTMSG_(
				ops_functors.find(selected_op) != ops_functors.end(),
				"Internal error: Unknown operation functor!");

			// Call the selected functor:
			ops_functors[selected_op](input_file, is3d, cmd, verbose);
		}

		// successful end of program.
		ret_val = 0;
	}
	catch (const std::exception& e)
	{
		std::cerr << mrpt::exception_to_str(e) << std::endl;
		ret_val = -1;
	}

	// end:
	return ret_val;
}

// ===================================================================================
//     Implementation of the functions operations
// ===================================================================================

// -----------------------------------------------------------------------------------
//  op: --view
//   Load and visualize the graph
// -----------------------------------------------------------------------------------
IMPLEMENT_OP_FUNCTION(op_view)
{
	// Load:
	GRAPHTYPE g;
	g.loadFromTextFile(in_file);

	// Display:
	display_graph(g);
}

// -----------------------------------------------------------------------------------
//  op: --info
//   Load and visualize basic info on graph
// -----------------------------------------------------------------------------------
IMPLEMENT_OP_FUNCTION(op_info)
{
	// Load:
	GRAPHTYPE g;
	g.loadFromTextFile(in_file);

	set<TNodeID> lstNodeIDs;
	g.getAllNodes(lstNodeIDs);

	// Show graph stats:
	cout << "Edge count                         : " << g.edgeCount() << endl;
	cout << "Nodes count (in VERTEX2/3 entries) : " << g.nodes.size() << endl;
	cout << "Nodes count (in edge entries)      : " << lstNodeIDs.size()
		 << endl;
}

// -----------------------------------------------------------------------------------
//  op: --dijkstra
//   Load and visualize the graph
// -----------------------------------------------------------------------------------
IMPLEMENT_OP_FUNCTION(op_dijkstra)
{
	const bool save_to_file = arg_output_file.isSet();	// Output to file??
	const bool display_3D = arg_view.isSet();  // Output to 3D view??

	if (!save_to_file && !display_3D)
		std::cerr << "\n ** WARNING **: Neither --view nor --output specified. "
					 "The result will not be saved anywhere.\n\n";

	// Load:
	GRAPHTYPE g;
	g.loadFromTextFile(in_file);

	// Find the first node ID and use it as root:
	VERBOSE_COUT << "Making a list with all node IDs...\n";

	set<TNodeID> lstNodeIDs;
	g.getAllNodes(lstNodeIDs);

	ASSERT_(!lstNodeIDs.empty());
	const TNodeID id_root = *lstNodeIDs.begin();
	VERBOSE_COUT << "Using root node ID=" << id_root << endl;

	g.root = id_root;

	VERBOSE_COUT << "Executing Dijkstra...\n";
	CTicTac tictac;
	tictac.Tic();

	g.dijkstra_nodes_estimate();

	VERBOSE_COUT << "Took: " << mrpt::system::formatTimeInterval(tictac.Tac())
				 << endl;

	if (save_to_file)
	{
		const string out_file = arg_output_file.getValue();
		VERBOSE_COUT << "Saving resulting graph to: " << out_file << endl;
		g.saveToTextFile(out_file);
	}
	if (display_3D)
	{
		VERBOSE_COUT << "Displaying resulting graph\n";
		display_graph(g);
	}
}  // end op_dijkstra

// -----------------------------------------------------------------------------------
//  op: --levmarq
//   Load and optimize a graph with the sparse Lev-Marq algorithm.
// -----------------------------------------------------------------------------------
IMPLEMENT_OP_FUNCTION(op_levmarq)
{
	const bool save_to_file = arg_output_file.isSet();	// Output to file??
	const bool display_3D = arg_view.isSet();  // Output to 3D view??
	const bool skip_dijkstra = arg_no_span.isSet();

	if (!save_to_file && !display_3D)
		std::cerr << "\n ** WARNING **: Neither --view nor --output specified. "
					 "The result will not be saved anywhere.\n\n";

	// Load:
	GRAPHTYPE g;
	g.loadFromTextFile(in_file);

	// Find the first node ID and use it as root:
	VERBOSE_COUT << "Making a list with all node IDs...\n";

	set<TNodeID> lstNodeIDs;
	g.getAllNodes(lstNodeIDs);

	ASSERT_(!lstNodeIDs.empty());
	const TNodeID id_root = *lstNodeIDs.begin();
	VERBOSE_COUT << "Using root node ID=" << id_root << endl;

	g.root = id_root;

	// Executes dijkstra (spanning tree) for initial guess?
	// --------------------------------
	if (!skip_dijkstra)
	{
		VERBOSE_COUT << "Executing Dijkstra...\n";
		CTicTac tictac;
		tictac.Tic();

		g.dijkstra_nodes_estimate();

		VERBOSE_COUT << "Took: "
					 << mrpt::system::formatTimeInterval(tictac.Tac()) << endl;
	}

	// Executes the optimization:
	// --------------------------------
	mrpt::containers::yaml params;
	params["verbose"] = verbose;
	params["profiler"] = verbose;
	params["max_iterations"] = arg_max_iters.getValue();
	params["initial_lambda"] = arg_initial_lambda.getValue();

	graphslam::TResultInfoSpaLevMarq info;

	graphslam::optimize_graph_spa_levmarq(g, info, nullptr, params);

	// Output:
	if (save_to_file)
	{
		const string out_file = arg_output_file.getValue();
		VERBOSE_COUT << "Saving resulting graph to: " << out_file << endl;
		g.saveToTextFile(out_file);
	}
	if (display_3D)
	{
		VERBOSE_COUT << "Displaying resulting graph\n";
		display_graph(g);
	}
}  // end op_levmarq
