/* +------------------------------------------------------------------------+
   |                     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 "nav-precomp.h"  // Precomp header
//
#include <mrpt/nav/planners/PlannerRRT_SE2_TPS.h>
#include <mrpt/nav/tpspace/CPTG_DiffDrive_CollisionGridBased.h>
#include <mrpt/random.h>
#include <mrpt/system/CTicTac.h>
#include <mrpt/system/filesystem.h>

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

MRPT_TODO("Optimize getNearestNode() with KD-tree!")

PlannerRRT_SE2_TPS::PlannerRRT_SE2_TPS() = default;
/** Load all params from a config file source */
void PlannerRRT_SE2_TPS::loadConfig(
	const mrpt::config::CConfigFileBase& ini, const std::string& sSect)
{
	PlannerTPS_VirtualBase::internal_loadConfig_PTG(ini, sSect);
}

/** Must be called after setting all params (see `loadConfig()`) and before
 * calling `solve()` */
void PlannerRRT_SE2_TPS::initialize()
{
	PlannerTPS_VirtualBase::internal_initialize_PTG();

	m_initialized = true;
}

/** The main API entry point: tries to find a planned path from 'goal' to
 * 'target' */
void PlannerRRT_SE2_TPS::solve(
	const PlannerRRT_SE2_TPS::TPlannerInput& pi,
	PlannerRRT_SE2_TPS::TPlannerResult& result)
{
	mrpt::system::CTimeLoggerEntry tleg(m_timelogger, "PT_RRT::solve");

	// Sanity checks:
	ASSERTMSG_(m_initialized, "initialize() must be called before!");

	// Calc maximum vehicle shape radius:
	double max_veh_radius = 0.;
	for (const auto& ptg : m_PTGs)
		mrpt::keep_max(max_veh_radius, ptg->getMaxRobotRadius());

	// [Algo `tp_space_rrt`: Line 1]: Init tree adding the initial pose
	if (result.move_tree.getAllNodes().empty())
	{
		result.move_tree.root = 0;
		result.move_tree.insertNode(
			result.move_tree.root, TNodeSE2_TP(pi.start_pose));
	}

	mrpt::system::CTicTac working_time;
	working_time.Tic();
	size_t rrt_iter_counter = 0;

	size_t SAVE_3D_TREE_LOG_DECIMATION_CNT = 0;
	static size_t SAVE_LOG_SOLVE_COUNT = 0;
	SAVE_LOG_SOLVE_COUNT++;

	// Keep track of the best solution so far:
	// By reusing the contents of "result" we make the algorithm re-callable
	// ("any-time" algorithm) to refine results

	// [Algo `tp_space_rrt`: Line 2]: Iterate
	// ------------------------------------------
	for (;;)
	{
		// Check end conditions:
		const double elap_tim = working_time.Tac();
		if ((end_criteria.maxComputationTime > 0 &&
			 elap_tim > end_criteria.maxComputationTime)  // Max comp time
			||
			(result.goal_distance < end_criteria.acceptedDistToTarget &&
			 elap_tim >= end_criteria.minComputationTime)  // Reach closer than
			// this to target
		)
		{
			break;
		}

		// [Algo `tp_space_rrt`: Line 3]: sample random state (with goal
		// biasing)
		// -----------------------------------------
		node_pose_t x_rand;
		// bool rand_is_target=false;
		if (mrpt::random::getRandomGenerator().drawUniform(0.0, 1.0) <
			params.goalBias)
		{
			x_rand = pi.goal_pose;
			// rand_is_target=true;
		}
		else
		{
			// Sample uniform:
			for (int i = 0; i < node_pose_t::static_size; i++)
				x_rand[i] = mrpt::random::getRandomGenerator().drawUniform(
					pi.world_bbox_min[i], pi.world_bbox_max[i]);
		}
		const CPose2D x_rand_pose(x_rand);

		// [Algo `tp_space_rrt`: Line 4]: Init empty solution set
		// -----------------------------------------
		using sorted_solution_list_t = std::map<double, TMoveEdgeSE2_TP>;
		sorted_solution_list_t candidate_new_nodes;	 // Map: cost -> info. Pick
		// begin() to select the
		// lowest-cose one.

		const PoseDistanceMetric<TNodeSE2>
			distance_evaluator_se2;	 // Plain distances in SE(2), not along PTGs
		bool is_new_best_solution = false;	// Just for logging purposes

		//#define DO_LOG_TXTS
		std::string sLogTxt;

		// [Algo `tp_space_rrt`: Line 5]: For each PTG
		// -----------------------------------------
		const size_t nPTGs = m_PTGs.size();
		for (size_t idxPTG = 0; idxPTG < nPTGs; ++idxPTG)
		{
			rrt_iter_counter++;

			// [Algo `tp_space_rrt`: Line 5]: Search nearest neig. to x_rand
			// -----------------------------------------------
			const PoseDistanceMetric<TNodeSE2_TP> distance_evaluator(
				*m_PTGs[idxPTG]);

			const TNodeSE2_TP query_node(x_rand);

			m_timelogger.enter("TMoveTree::getNearestNode");
			mrpt::graphs::TNodeID x_nearest_id =
				result.move_tree.getNearestNode(query_node, distance_evaluator);
			m_timelogger.leave("TMoveTree::getNearestNode");

			if (x_nearest_id == mrpt::graphs::INVALID_NODEID)
			{
				// We can't find any close node, at least with this PTG's paths:
				// skip

				// Before that, save log:
				if (params.save_3d_log_freq > 0 &&
					(++SAVE_3D_TREE_LOG_DECIMATION_CNT >=
					 params.save_3d_log_freq))
				{
					SAVE_3D_TREE_LOG_DECIMATION_CNT =
						0;	// Reset decimation counter
					TRenderPlannedPathOptions render_options;
					render_options.highlight_path_to_node_id =
						result.best_goal_node_id;
					render_options.highlight_last_added_edge = false;
					render_options.x_rand_pose = &x_rand_pose;
					render_options.log_msg = "SKIP: Can't find any close node";
					render_options.log_msg_position = mrpt::math::TPoint3D(
						pi.world_bbox_min.x, pi.world_bbox_min.y, 0);
					render_options.ground_xy_grid_frequency = 1.0;

					mrpt::opengl::Scene scene;
					renderMoveTree(scene, pi, result, render_options);
					mrpt::system::createDirectory("./rrt_log_trees");
					scene.saveToFile(mrpt::format(
						"./rrt_log_trees/rrt_log_%03u_%06u.3Dscene",
						static_cast<unsigned int>(SAVE_LOG_SOLVE_COUNT),
						static_cast<unsigned int>(rrt_iter_counter)));
				}

				continue;  // Skip
			}

			const TNodeSE2_TP& x_nearest_node =
				result.move_tree.getAllNodes().find(x_nearest_id)->second;

			// [Algo `tp_space_rrt`: Line 6]: Relative target
			// -----------------------------------------------
			const CPose2D x_nearest_pose(x_nearest_node.state);
			const CPose2D x_rand_rel = x_rand_pose - x_nearest_pose;

			// [Algo `tp_space_rrt`: Line 7]: Relative target in TP-Space
			// ------------------------------------------------------------
			const double D_max =
				std::min(params.maxLength, m_PTGs[idxPTG]->getRefDistance());

			double d_rand;	// Coordinates in TP-space
			int k_rand;	 // k_rand is the index of target_alpha in PTGs
			// corresponding to a specific d_rand
			// bool tp_point_is_exact =
			m_PTGs[idxPTG]->inverseMap_WS2TP(
				x_rand_rel.x(), x_rand_rel.y(), k_rand, d_rand);
			d_rand *=
				m_PTGs[idxPTG]
					->getRefDistance();	 // distance to target, in "real meters"

			float d_free;
			// bool local_obs_ok = false; // Just for 3D log files: indicates
			// whether obstacle points have been recomputed

			// [Algo `tp_space_rrt`: Line 8]: TP-Obstacles
			// ------------------------------------------------------------
			// Transform obstacles as seen from x_nearest_node -> TP_obstacles
			double TP_Obstacles_k_rand = .0;  // vector<double> TP_Obstacles;
			const double MAX_DIST_FOR_OBSTACLES =
				1.5 * m_PTGs[idxPTG]->getRefDistance();	 // Maximum Euclidean
			// distance (radius)
			// for considering
			// obstacles around the
			// current robot pose

			ASSERT_GT_(
				m_PTGs[idxPTG]->getRefDistance(),
				1.1 * max_veh_radius);	// Make sure the PTG covers at least a
			// bit more than the vehicle shape!!
			// (should be much, much higher)

			{
				CTimeLoggerEntry tle(
					m_timelogger, "PT_RRT::solve.changeCoordinatesReference");
				transformPointcloudWithSquareClipping(
					pi.obstacles_points, m_local_obs,
					CPose2D(x_nearest_node.state), MAX_DIST_FOR_OBSTACLES);
				// local_obs_ok=true;
			}
			{
				CTimeLoggerEntry tle(
					m_timelogger, "PT_RRT::solve.SpaceTransformer");
				spaceTransformerOneDirectionOnly(
					k_rand, m_local_obs, m_PTGs[idxPTG].get(),
					MAX_DIST_FOR_OBSTACLES, TP_Obstacles_k_rand);
			}

			// directions k_rand in TP_obstacles[k_rand] = d_free
			// this is the collision free distance to the TP_target
			d_free = TP_Obstacles_k_rand;  // TP_Obstacles[k_rand];

			// [Algo `tp_space_rrt`: Line 10]: d_new
			// ------------------------------------------------------------
			double d_new = std::min(
				D_max,
				d_rand);  // distance of the new candidate state in TP-space

			// mrpt::poses::CPose2D *log_new_state_ptr=NULL; // For graphical
			// logs only

#ifdef DO_LOG_TXTS
			sLogTxt += mrpt::format(
				"tp_idx=%u tp_exact=%c\n d_free: %f d_rand=%f d_new=%f\n",
				static_cast<unsigned int>(idxPTG),
				tp_point_is_exact ? 'Y' : 'N', d_free, d_rand, d_new);
			sLogTxt += mrpt::format(
				" nearest:%s\n", x_nearest_pose.asString().c_str());
#endif

			// [Algo `tp_space_rrt`: Line 13]: Do we have free space?
			// ------------------------------------------------------------
			if (d_free >= d_new)
			{
				// [Algo `tp_space_rrt`: Line 14]: PTG function
				// ------------------------------------------------------------
				// given d_rand and k_rand provides x,y,phi of the point in
				// c-space
				uint32_t nStep;
				m_PTGs[idxPTG]->getPathStepForDist(k_rand, d_new, nStep);

				mrpt::math::TPose2D rel_pose =
					m_PTGs[idxPTG]->getPathPose(k_rand, nStep);

				// wrap to [-pi,pi] to avoid out of bounds errors:
				mrpt::math::wrapToPiInPlace(rel_pose.phi);

				// [Algo `tp_space_rrt`: Line 15]: pose composition
				// ------------------------------------------------------------
				const mrpt::poses::CPose2D new_state_rel(rel_pose);
				mrpt::poses::CPose2D new_state =
					x_nearest_pose + new_state_rel;	 // compose the new_motion
				// as the last nmotion and
				// the new state
				// log_new_state_ptr = &new_state;

				// Check whether there's already a too-close node around:
				// --------------------------------------------------------
				bool accept_this_node = true;

				// Is this a potential solution
				const double goal_dist =
					new_state.distance2DTo(pi.goal_pose.x, pi.goal_pose.y);
				const double goal_ang = std::abs(
					mrpt::math::angDistance(new_state.phi(), pi.goal_pose.phi));
				const bool is_acceptable_goal =
					(goal_dist < end_criteria.acceptedDistToTarget) &&
					(goal_ang < end_criteria.acceptedAngToTarget);

				auto new_nearest_id = mrpt::graphs::INVALID_NODEID;
				if (!is_acceptable_goal)  // Only check for nearby nodes if this
				// is not a solution!
				{
					double new_nearest_dist;
					const TNodeSE2 new_state_node(new_state.asTPose());

					m_timelogger.enter("TMoveTree::getNearestNode");
					new_nearest_id = result.move_tree.getNearestNode(
						new_state_node, distance_evaluator_se2,
						&new_nearest_dist, &result.acceptable_goal_node_ids);
					m_timelogger.leave("TMoveTree::getNearestNode");

					if (new_nearest_id != mrpt::graphs::INVALID_NODEID)
					{
						// Also check angular distance:
						const double new_nearest_ang =
							std::abs(mrpt::math::angDistance(
								new_state.phi(),
								result.move_tree.getAllNodes()
									.find(new_nearest_id)
									->second.state.phi));
						accept_this_node =
							(new_nearest_dist >=
								 params.minDistanceBetweenNewNodes ||
							 new_nearest_ang >= params.minAngBetweenNewNodes);
					}
				}

				if (!accept_this_node)
				{
#ifdef DO_LOG_TXTS
					if (new_nearest_id != mrpt::graphs::INVALID_NODEID)
					{
						sLogTxt += mrpt::format(
							" -> new node NOT accepted for closeness to: %s\n",
							result.move_tree.getAllNodes()
								.find(new_nearest_id)
								->second.state.asString()
								.c_str());
					}
#endif
					continue;  // Too close node, skip!
				}

				// [Algo `tp_space_rrt`: Line 16]: Add to candidate solution set
				// ------------------------------------------------------------
				// Create "movement" (tree edge) object:
				TMoveEdgeSE2_TP new_edge(x_nearest_id, new_state.asTPose());

				new_edge.cost = d_new;
				new_edge.ptg_index = idxPTG;
				new_edge.ptg_K = k_rand;
				new_edge.ptg_dist = d_new;

				candidate_new_nodes[new_edge.cost] = new_edge;

			}  // end if the path is obstacle free
			else
			{
#ifdef DO_LOG_TXTS
				sLogTxt += mrpt::format(" -> d_free NOT < d_rand\n");
#endif
			}

		}  // end for idxPTG

		// [Algo `tp_space_rrt`: Line 19]: Any solution found?
		// ------------------------------------------------------------
		if (!candidate_new_nodes.empty())
		{
			const TMoveEdgeSE2_TP& best_edge =
				candidate_new_nodes.begin()->second;
			const TNodeSE2_TP new_state_node(best_edge.end_state);

			// Insert into the tree:
			const mrpt::graphs::TNodeID new_child_id =
				result.move_tree.getNextFreeNodeID();
			result.move_tree.insertNodeAndEdge(
				best_edge.parent_id, new_child_id, new_state_node, best_edge);

			// Distance to goal:
			const double goal_dist =
				mrpt::poses::CPose2D(best_edge.end_state)
					.distance2DTo(pi.goal_pose.x, pi.goal_pose.y);
			const double goal_ang = std::abs(mrpt::math::angDistance(
				best_edge.end_state.phi, pi.goal_pose.phi));

			const bool is_acceptable_goal =
				(goal_dist < end_criteria.acceptedDistToTarget) &&
				(goal_ang < end_criteria.acceptedAngToTarget);

			if (is_acceptable_goal)
				result.acceptable_goal_node_ids.insert(new_child_id);

			// Total path length:
			double this_path_cost = std::numeric_limits<double>::max();
			if (is_acceptable_goal)	 // Don't waste time computing path length
			// if it doesn't matter anyway
			{
				TMoveTreeSE2_TP::path_t candidate_solution_path;
				result.move_tree.backtrackPath(
					new_child_id, candidate_solution_path);
				this_path_cost = 0;
				for (auto it = candidate_solution_path.begin();
					 it != candidate_solution_path.end(); ++it)
					if (it->edge_to_parent)
						this_path_cost += it->edge_to_parent->cost;
			}

			// Check if this should be the new optimal path:
			if (is_acceptable_goal && this_path_cost < result.path_cost)
			{
				result.goal_distance = goal_dist;
				result.path_cost = this_path_cost;

				result.best_goal_node_id = new_child_id;
				is_new_best_solution = true;
			}
		}  // end if any candidate found

		//  Graphical logging, if enabled:
		// ------------------------------------------------------
		if (params.save_3d_log_freq > 0 &&
			(++SAVE_3D_TREE_LOG_DECIMATION_CNT >= params.save_3d_log_freq ||
			 is_new_best_solution))
		{
			CTimeLoggerEntry tle(
				m_timelogger, "PT_RRT::solve.generate_log_files");
			SAVE_3D_TREE_LOG_DECIMATION_CNT = 0;  // Reset decimation counter

			// Render & save to file:
			TRenderPlannedPathOptions render_options;
			render_options.highlight_path_to_node_id = result.best_goal_node_id;
			render_options.x_rand_pose = &x_rand_pose;
			// render_options.x_nearest_pose = &x_nearest_pose;
			// if (local_obs_ok) render_options.local_obs_from_nearest_pose =
			// &m_local_obs;
			// render_options.new_state = log_new_state_ptr;
			render_options.highlight_last_added_edge = true;
			render_options.ground_xy_grid_frequency = 1.0;

			render_options.log_msg = sLogTxt;
			render_options.log_msg_position = mrpt::math::TPoint3D(
				pi.world_bbox_min.x, pi.world_bbox_min.y, 0);

			mrpt::opengl::Scene scene;
			renderMoveTree(scene, pi, result, render_options);

			mrpt::system::createDirectory("./rrt_log_trees");
			scene.saveToFile(mrpt::format(
				"./rrt_log_trees/rrt_log_%03u_%06u.3Dscene",
				static_cast<unsigned int>(SAVE_LOG_SOLVE_COUNT),
				static_cast<unsigned int>(rrt_iter_counter)));
		}

	}  // end loop until end conditions

	// [Algo `tp_space_rrt`: Line 17]: Tree back trace
	// ------------------------------------------------------------
	result.success = (result.goal_distance < end_criteria.acceptedDistToTarget);
	result.computation_time = working_time.Tac();

}  // end solve()
