#include "MultilevelAStarEx.h"

#include <godot_cpp/core/class_db.hpp>
#include <godot_cpp/core/error_macros.hpp>
#include <godot_cpp/variant/utility_functions.hpp>

#include <set>

#define TERRAIN(x, y) (_terrain[(x) + (y) * _width])
#define UNITS(x, y) (_units[(x) + (y) * _width])
#define NODES(x, y) (_nodes[(x) + (y) * _width])

using namespace godot;

static const int STRAIGHT_DISTANCE = 10;
static const int DIAGONAL_DISTANCE = 14;

Node::Node()
{
	this->openPass = 0;
	this->closedPass = 0;
}

Node::Node(int x, int y) : Node()
{
	this->x = x;
	this->y = y;
}

void Node::open(int pass, Node *parent, int distanceFromStart, const Vector2i &end)
{
	DEV_ASSERT(this->openPass != pass);

	this->openPass = pass;
	this->parent = parent;
	this->distanceFromStart = distanceFromStart;

	int dx = abs(x - end.x);
	int dy = abs(y - end.y);

	int max, min;
	if (dx > dy)
	{
		max = dx;
		min = dy;
	}
	else
	{
		max = dy;
		min = dx;
	}
	this->distanceToEnd = min * 14 + (max - min) * 10;

	if ((dx == 0 || dx == 1) && (dy == 0 || dy == 1))
	{
		this->distanceToEndForClosest = STRAIGHT_DISTANCE;
	}
	else
	{
		this->distanceToEndForClosest = this->distanceToEnd;
	}
}

void Node::close(int pass)
{
	DEV_ASSERT(openPass == pass && closedPass != pass);

	closedPass = pass;
}

Node::NodeState Node::state(int pass) const
{
	if (closedPass == pass)
	{
		return CLOSED;
	}
	else if (openPass == pass)
	{
		return OPEN;
	}
	else
	{
		return UNUSED;
	}
}

int Node::total_cost() const
{
	return distanceFromStart + distanceToEnd;
}

bool Node::compare_nodes(const Node *left, const Node *right) {
	if (left->total_cost() < right->total_cost())
	{
		return true;
	}
	else if (left->total_cost() > right->total_cost())
	{
		return false;
	}
	else // if (left->total_cost() == right->total_cost())
	{
		if (left->distanceToEnd < right->distanceToEnd)
		{
			return true;
		}
		else if (left->distanceToEnd > right->distanceToEnd)
		{
			return false;
		}
		else // if (left->distanceToEnd == right->distanceToEnd)
		{
			// make sure no two nodes are the same
			return left < right;
		}
	}
}

void MultilevelAStarEx::_bind_methods()
{
	ClassDB::bind_method(D_METHOD("init", "region"), &MultilevelAStarEx::init);
	ClassDB::bind_method(D_METHOD("get_region"), &MultilevelAStarEx::get_region);
	ClassDB::bind_method(D_METHOD("set_terrain", "cell", "type"), &MultilevelAStarEx::set_terrain);
	ClassDB::bind_method(D_METHOD("get_terrain", "cell"), &MultilevelAStarEx::get_terrain);
	ClassDB::bind_method(D_METHOD("set_unit", "cell", "blocked"), &MultilevelAStarEx::set_unit);
	ClassDB::bind_method(D_METHOD("get_unit", "cell"), &MultilevelAStarEx::get_unit);
	ClassDB::bind_method(D_METHOD("find_path", "from", "to", "return_closest"), &MultilevelAStarEx::find_path);

	BIND_ENUM_CONSTANT(STAIRS);
	BIND_ENUM_CONSTANT(BLOCKED);
	BIND_ENUM_CONSTANT(GROUND);
}

MultilevelAStarEx::MultilevelAStarEx()
{
	//UtilityFunctions::print("Constructor.");

	_init = false;
	_pass = 0;
}

MultilevelAStarEx::~MultilevelAStarEx() { }

void MultilevelAStarEx::init(const Rect2i &region)
{
	DEV_ASSERT(!_init);

	DEV_ASSERT(region.get_area() >= 0);

	_region = region;
	_width = region.get_size().width;
	_height = region.get_size().height;
	_trans = region.get_position();

	_terrain.resize(_width * _height, BLOCKED);
	_units.resize(_width * _height, false);
	_nodes.resize(_width * _height);

	std::vector<Node>::iterator iter = _nodes.begin();
	for (int y = 0; y < _height; y++)
	{
		for (int x = 0; x < _width; x++)
		{
			*iter++ = Node(x, y);
		}
	}

	_init = true;
}

Rect2i MultilevelAStarEx::get_region() const
{
	DEV_ASSERT(_init);

	return _region;
}

void MultilevelAStarEx::set_terrain(const Vector2i &cell, TerrainType type)
{
	DEV_ASSERT(_init);
	DEV_ASSERT(_region.has_point(cell));

	Vector2i cell2 = cell - _trans;
	TERRAIN(cell2.x, cell2.y) = type;
}

MultilevelAStarEx::TerrainType MultilevelAStarEx::get_terrain(const Vector2i &cell) const
{
	DEV_ASSERT(_init);
	DEV_ASSERT(_region.has_point(cell));

	Vector2i cell2 = cell - _trans;
	return TERRAIN(cell2.x, cell2.y);
}

void MultilevelAStarEx::set_unit(const Vector2i &cell, bool blocked)
{
	DEV_ASSERT(_init);
	DEV_ASSERT(_region.has_point(cell));

	Vector2i cell2 = cell - _trans;
	UNITS(cell2.x, cell2.y) = blocked;
}

bool MultilevelAStarEx::get_unit(const Vector2i &cell) const
{
	DEV_ASSERT(_init);
	DEV_ASSERT(_region.has_point(cell));

	Vector2i cell2 = cell - _trans;
	return UNITS(cell2.x, cell2.y);
}

bool MultilevelAStarEx::can_move(const Node *current, int x, int y) const
{
	if (UNITS(x, y)) return false;

	TerrainType tc = TERRAIN(current->x, current->y);
	TerrainType td = TERRAIN(x, y);

	if (td == BLOCKED)
	{
		return false;
	}
	else if (td == tc)
	{
		return true;
	}
	else if (td < 0 && tc > 0) // td is stairs
	{
		if (-td == tc || -td == tc - 1) return true;
	}
	else if (td > 0 && tc < 0) // tc is stairs
	{
		if (-tc == td || -tc == td - 1) return true;
	}
	else if (td < 0 && tc < 0) // both are stairs
	{
		return abs(td - tc) < 2;
	}

	return false;
}

TypedArray<Vector2i> MultilevelAStarEx::generate_path(const Node *current) const
{
	// count the number of nodes in the path
	int cnt = 0;
	for (const Node *n = current; n != nullptr; n = n->parent) cnt++;

	TypedArray<Vector2i> arr;
	arr.resize(cnt);

	int i = cnt;
	while (current != nullptr)
	{
		arr[--cnt] = Vector2i(current->x, current->y) + _trans;
		current = current->parent;
	}

	return arr;
}

TypedArray<Vector2i> MultilevelAStarEx::find_path(const Vector2i &from, const Vector2i &to, bool return_closest)
{
	DEV_ASSERT(_init);

	DEV_ASSERT(_region.has_point(from));
	DEV_ASSERT(_region.has_point(to));

	_pass++;

	Vector2i from2 = from - _trans;
	Vector2i to2 = to - _trans;

	std::set<Node *, decltype(Node::compare_nodes)*> open(&Node::compare_nodes);

	Node *closest = &NODES(from2.x, from2.y);
	closest->open(_pass, nullptr, 0, to2);
	open.insert(closest);

	auto process = [this, &open, &to2, &closest](Node *current, int x, int y, int distance) {
		Node *node = &NODES(x, y);
		
		if (node->state(_pass) == Node::CLOSED)
		{
			return;
		}
		else if (node->state(_pass) == Node::UNUSED)
		{
			node->open(_pass, current, current->distanceFromStart + distance, to2);
			open.insert(node);
		}
		else if (current->distanceFromStart + distance < node->distanceFromStart)
		{
			auto nodeIter = open.find(node);
			DEV_ASSERT(nodeIter != open.end());
			open.erase(nodeIter);

			node->parent = current;
			node->distanceFromStart = current->distanceFromStart + distance;

			open.insert(node);
		}

		// find the closest cell
		if (node->distanceToEndForClosest < closest->distanceToEndForClosest)
		{
			closest = node;
		}
		else if (node->distanceToEndForClosest == closest->distanceToEndForClosest)
		{
			if (node->distanceFromStart < closest->distanceFromStart)
			{
				closest = node;
			}
		}
	};

	while (!open.empty())
	{
		Node *current = *open.begin();

		if (current->distanceToEnd == 0)
		{
			// found the path
			return generate_path(current);
		}

		// close it
		open.erase(open.begin());
		current->close(_pass);

		// expand it
		if (current->x - 1 >= 0) // left
		{
			if (can_move(current, current->x - 1, current->y))
			{
				process(current, current->x - 1, current->y, STRAIGHT_DISTANCE);
			}
		}
		if (current->x + 1 < _width) // right
		{
			if (can_move(current, current->x + 1, current->y))
			{
				process(current, current->x + 1, current->y, STRAIGHT_DISTANCE);
			}
		}
		if (current->y - 1 >= 0) // up
		{
			if (can_move(current, current->x, current->y - 1))
			{
				process(current, current->x, current->y - 1, STRAIGHT_DISTANCE);
			}
		}
		if (current->y + 1 < _height) // down
		{
			if (can_move(current, current->x, current->y + 1))
			{
				process(current, current->x, current->y + 1, STRAIGHT_DISTANCE);
			}
		}
		if ((current->x - 1 >= 0) && (current->y - 1 >= 0)) // top left
		{
			if (can_move(current, current->x - 1, current->y - 1))
			{
				process(current, current->x - 1, current->y - 1, DIAGONAL_DISTANCE);
			}
		}
		if ((current->x + 1 < _width) && (current->y - 1 >= 0)) // top right
		{
			if (can_move(current, current->x + 1, current->y - 1))
			{
				process(current, current->x + 1, current->y - 1, DIAGONAL_DISTANCE);
			}
		}
		if ((current->x - 1 >= 0) && (current->y + 1 < _height)) // bottom left
		{
			if (can_move(current, current->x - 1, current->y + 1))
			{
				process(current, current->x - 1, current->y + 1, DIAGONAL_DISTANCE);
			}
		}
		if ((current->x + 1 < _width) && (current->y + 1 < _height)) // bottom right
		{
			if (can_move(current, current->x + 1, current->y + 1))
			{
				process(current, current->x + 1, current->y + 1, DIAGONAL_DISTANCE);
			}
		}
	}

	if (return_closest)
	{
		// return path to closest
		return generate_path(closest);
	}

	return TypedArray<Vector2i>();
}