#pragma once

#include <psemek/util/functional.hpp>
#include <psemek/util/hash_table.hpp>

#include <set>
#include <deque>
#include <optional>

namespace psemek::util
{

	template <typename Cost, typename Node, typename NeighboursFn, typename HeuristicFn, typename EdgeCallback>
	struct pathfinder
	{
		struct node_compare
		{
			util::hash_map<Node, Cost> const & priority;

			bool operator()(Node const & n1, Node const & n2) const
			{
				auto const p1 = priority.at(n1);
				auto const p2 = priority.at(n2);

				if (p1 < p2)
					return true;
				if (p1 > p2)
					return false;
				return n1 < n2;
			}
		};

		NeighboursFn node_neighbours;
		HeuristicFn heuristic;
		EdgeCallback edge_callback;

		util::hash_map<Node, Cost> cost;
		util::hash_map<Node, Cost> priority;
		util::hash_map<Node, Node> previous;
		std::set<Node, node_compare> queue;

		template <typename NeighboursFn1, typename HeuristicFn1, typename EdgeCallback1>
		pathfinder(NeighboursFn1 && node_neighbours, HeuristicFn1 && heuristic, EdgeCallback1 && edge_callback)
			: node_neighbours(std::forward<NeighboursFn1>(node_neighbours))
			, heuristic(std::forward<HeuristicFn1>(heuristic))
			, edge_callback(std::forward<EdgeCallback1>(edge_callback))
			, queue(node_compare{priority})
		{}

		void reset()
		{
			queue.clear();
			previous.clear();
			priority.clear();
			cost.clear();
		}

		void init(Node const & start)
		{
			cost[start] = Cost{};
			priority[start] = heuristic(start);
			queue.insert(start);
		}

		bool finished() const
		{
			return queue.empty();
		}

		Node step()
		{
			Node const node = std::move(*queue.begin());
			queue.erase(queue.begin());

			auto const node_cost = cost[node];

			node_neighbours(node, [&]<typename EdgeId = int>(Node neighbour, Cost const & edge_cost, EdgeId const & edge_id = 0){
				Cost const new_cost = node_cost + edge_cost;

				auto it = cost.find(neighbour);

				if (it == cost.end() || new_cost < it->second)
				{
					edge_callback(node, node_cost, neighbour, edge_cost, edge_id, true);

					if (it != cost.end())
						queue.erase(neighbour);

					cost[neighbour] = new_cost;
					priority[neighbour] = new_cost + heuristic(neighbour);
					previous[neighbour] = node;
					queue.insert(std::move(neighbour));
				}
				else
				{
					edge_callback(node, node_cost, neighbour, edge_cost, edge_id, false);
				}
			});

			return node;
		}

		bool found(Node const & end) const
		{
			return cost.contains(end);
		}

		// Doesn't include the starting node
		template <typename Iterator, typename IsStart>
		Iterator path(IsStart const & is_start, Node const & end, Iterator it) const
		{
			Node node = end;
			while (true)
			{
				if (is_start(node))
					break;
				*it++ = node;
				node = previous.at(node);
			}
			return it;
		}

		// Includes the starting node
		template <typename Iterator, typename IsStart>
		Iterator path_full(IsStart const & is_start, Node const & end, Iterator it) const
		{
			Node node = end;
			while (true)
			{
				*it++ = node;
				if (is_start(node))
					break;
				node = previous.at(node);
			}
			return it;
		}

		// Doesn't include the starting node
		template <typename Iterator>
		Iterator path(Node const & start, Node const & end, Iterator it) const
		{
			return path([start](Node const & node){ return node == start; }, end, it);
		}

		// Doesn't include the starting node
		std::deque<Node> path(Node const & start, Node const & end) const
		{
			std::deque<Node> result;
			path(start, end, std::front_inserter(result));
			return result;
		}
	};

	template <typename Cost, typename Node, typename NeighboursFn, typename HeuristicFn, typename EdgeCallback>
	auto make_pathfinder(NeighboursFn && node_neighbours, HeuristicFn && heuristic, EdgeCallback && edge_callback)
	{
		return pathfinder<Cost, Node, NeighboursFn, HeuristicFn, EdgeCallback>(std::forward<NeighboursFn>(node_neighbours), std::forward<HeuristicFn>(heuristic), std::forward<EdgeCallback>(edge_callback));
	}

	template <typename Cost, typename Node, typename NeighboursFn, typename HeuristicFn>
	auto make_pathfinder(NeighboursFn && node_neighbours, HeuristicFn && heuristic)
	{
		return make_pathfinder<Cost, Node>(std::forward<NeighboursFn>(node_neighbours), std::forward<HeuristicFn>(heuristic), util::nop);
	}

	template <typename Cost, typename Node, typename NeighboursFn, typename HeuristicFn>
	std::optional<std::deque<Node>> find_path(Node const & start, Node const & end, NeighboursFn && node_neighbours, HeuristicFn && heuristic)
	{
		auto helper = make_pathfinder<Cost, Node>(std::forward<NeighboursFn>(node_neighbours), std::forward<HeuristicFn>(heuristic));

		helper.init(start);
		bool found = false;
		while (!helper.finished())
		{
			if (helper.step() == end)
			{
				found = true;
				break;
			}
		}

		if (found)
			return helper.path(start, end);

		return std::nullopt;
	}

}
