#include <psemek/ui/grid_layout.hpp>

#include <numeric>

namespace psemek::ui
{

	static float const inf = std::numeric_limits<float>::infinity();

	element::children_range grid_layout::children() const
	{
		return children_range{children_range_.data(), children_range_.data() + children_range_.size()};
	}

	bool grid_layout::add_child(std::shared_ptr<element> c)
	{
		for (std::size_t i = 0; i < row_count(); ++i)
		{
			for (std::size_t j = 0; j < column_count(); ++j)
			{
				if (!get(i, j))
				{
					set(i, j, std::move(c));
					return true;
				}
			}
		}
		return false;
	}

	bool grid_layout::has_child(element * c) const
	{
		for (std::size_t i = 0; i < row_count(); ++i)
			for (std::size_t j = 0; j < column_count(); ++j)
				if (get(i, j) == c) return true;
		return false;
	}

	std::shared_ptr<element> grid_layout::remove_child(element * c)
	{
		for (std::size_t i = 0; i < row_count(); ++i)
		{
			for (std::size_t j = 0; j < column_count(); ++j)
			{
				if (get(i, j) == c)
				{
					return remove(i, j);
				}
			}
		}
		return nullptr;
	}

	void grid_layout::set_row_count(std::size_t count)
	{
		set_size(count, std::max<std::size_t>(1, column_count()));
	}

	void grid_layout::set_column_count(std::size_t count)
	{
		set_size(std::max<std::size_t>(1, row_count()), count);
	}

	void grid_layout::set_size(std::size_t rows, std::size_t columns)
	{
		std::size_t const old_row_count = row_count();
		std::size_t const old_column_count = column_count();

		children_.resize({rows, columns});

		children_range_.resize(rows * columns);
		for (std::size_t i = 0; i < rows; ++i)
		{
			for (std::size_t j = 0; j < columns; ++j)
			{
				children_range_[i * columns + j] = children_(i, j).get();
			}
		}

		row_weight_.resize(rows);
		for (std::size_t i = old_row_count; i < rows; ++i)
			row_weight_[i] = 1.f;

		column_weight_.resize(columns);
		for (std::size_t i = old_column_count; i < columns; ++i)
			column_weight_[i] = 1.f;

		post_reshape();
	}

	void grid_layout::set_row_weight(std::size_t i, float w)
	{
		row_weight_.at(i) = w;
		post_reshape();
	}

	void grid_layout::set_column_weight(std::size_t i, float w)
	{
		column_weight_.at(i) = w;
		post_reshape();
	}

	element * grid_layout::get(std::size_t i, std::size_t j) const
	{
		return children_(i, j).get();
	}

	std::shared_ptr<element> grid_layout::set(std::size_t i, std::size_t j, std::shared_ptr<element> c)
	{
		auto old = std::move(children_(i, j));
		if (old) old->set_parent(nullptr);
		if (c) c->set_parent(this);
		children_range_[i * column_count() + j] = c.get();
		children_(i, j) = std::move(c);
		post_reshape();
		return old;
	}

	std::shared_ptr<element> grid_layout::remove(std::size_t i, std::size_t j)
	{
		return set(i, j, nullptr);
	}

	void grid_layout::set_outer_margin(bool value)
	{
		outer_margin_ = value;
		post_reshape();
	}

	static std::pair<std::vector<geom::interval<float>>, std::vector<geom::interval<float>>> row_column_sizes(util::array<std::shared_ptr<element>, 2> const & children)
	{
		std::vector<geom::interval<float>> row_size(children.width(), {0.f, inf});
		std::vector<geom::interval<float>> column_size(children.height(), {0.f, inf});

		for (std::size_t i = 0; i < children.width(); ++i)
		{
			for (std::size_t j = 0; j < children.height(); ++j)
			{
				auto c = children(i, j);
				if (!c) continue;

				auto sc = c->size_constraints();
				row_size[i] &= sc[1];
				column_size[j] &= sc[0];
			}
		}

		return {std::move(row_size), std::move(column_size)};
	}

	geom::box<float, 2> grid_layout::size_constraints() const
	{
		auto st = merged_style();
		if (!st) return element::size_constraints();

		auto row_column_size = row_column_sizes(children_);
		auto const & row_size = row_column_size.first;
		auto const & column_size = row_column_size.second;

		geom::box<float, 2> result{{{0.f, 0.f}, {0.f, 0.f}}};

		for (std::size_t i = 0; i < row_count(); ++i)
		{
			bool const minimized = (row_weight_[i] == 0.f);

			result[1].min += row_size[i].min;
			result[1].max += minimized ? row_size[i].min : row_size[i].max;
		}

		for (std::size_t i = 0; i < column_count(); ++i)
		{
			bool const minimized = (column_weight_[i] == 0.f);

			result[0].min += column_size[i].min;
			result[0].max += minimized ? column_size[i].min : column_size[i].max;
		}

		int const margin_row_count = (std::max<int>(1, row_count()) + (outer_margin_ ? 1 : -1));
		int const margin_column_count = (std::max<int>(1, column_count()) + (outer_margin_ ? 1 : -1));

		result[1] += margin_row_count * (*st->outer_margin);
		result[0] += margin_column_count * (*st->outer_margin);

		return result;
	}

	void grid_layout::reshape(geom::box<float, 2> const & bbox)
	{
		auto st = merged_style();
		if (!st) return;

		float const margin = *st->outer_margin;

		auto row_column_size = row_column_sizes(children_);
		auto const & row_size = row_column_size.first;
		auto const & column_size = row_column_size.second;

		int const margin_row_count = (std::max<int>(1, row_count()) + (outer_margin_ ? 1 : -1));
		int const margin_column_count = (std::max<int>(1, column_count()) + (outer_margin_ ? 1 : -1));

		float available_height = std::max(0.f, bbox[1].length() - margin * margin_row_count);
		float available_width = std::max(0.f, bbox[0].length() - margin * margin_column_count);

		float const row_weight_sum = std::accumulate(row_weight_.begin(), row_weight_.end(), 0.f);
		float const column_weight_sum = std::accumulate(column_weight_.begin(), column_weight_.end(), 0.f);

		std::vector<geom::interval<float>> row_shape(row_count());
		std::vector<geom::interval<float>> column_shape(column_count());

		// First, allocate minimized rows
		for (std::size_t i = 0; i < row_count(); ++i)
		{
			if (row_weight_[i] == 0.f)
			{
				row_shape[i] = {0.f, row_size[i].min};
				available_height -= row_size[i].min;
			}
		}

		available_height = std::max(available_height, 0.f);

		// Next, allocate scalable rows
		for (std::size_t i = 0; i < row_count(); ++i)
		{
			if (row_weight_[i] == 0.f)
				continue;

			float const row_height = available_height * row_weight_[i] / row_weight_sum;
			row_shape[i] = {0.f, row_height};
		}

		// Finally, move rows one after another
		for (std::size_t i = 0; i < row_count(); ++i)
		{
			if (i == 0)
				row_shape[i] += bbox[1].min + (outer_margin_ ? margin : 0.f);
			else
				row_shape[i] += row_shape[i - 1].max + margin;
		}

		// First, allocate minimized columns
		for (std::size_t i = 0; i < column_count(); ++i)
		{
			if (column_weight_[i] == 0.f)
			{
				column_shape[i] = {0.f, column_size[i].min};
				available_width -= column_size[i].min;
			}
		}

		available_width = std::max(available_width, 0.f);

		// Next, allocate scalable columns
		for (std::size_t i = 0; i < column_count(); ++i)
		{
			if (column_weight_[i] == 0.f)
				continue;

			float const column_width = available_width * column_weight_[i] / column_weight_sum;
			column_shape[i] = {0.f, column_width};
		}

		// Finally, move rows one after another
		for (std::size_t i = 0; i < column_count(); ++i)
		{
			if (i == 0)
				column_shape[i] += bbox[0].min + (outer_margin_ ? margin : 0.f);
			else
				column_shape[i] += column_shape[i - 1].max + margin;
		}

		// Apply calculated shapes
		for (std::size_t i = 0; i < row_count(); ++i)
		{
			for (std::size_t j = 0; j < column_count(); ++j)
			{
				if (!get(i, j)) continue;

				get(i, j)->reshape({{column_shape[j], row_shape[i]}});
			}
		}
	}

	grid_layout::~grid_layout()
	{
		release_children();
	}

}