#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).get() == 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).get() == 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();

		for (std::size_t i = 0; i < old_row_count; ++i)
		{
			for (std::size_t j = 0; j < old_column_count; ++j)
			{
				if (i >= rows || j >= columns)
					set(i, j, nullptr);
			}
		}

		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();
	}

	std::shared_ptr<element> grid_layout::get(std::size_t i, std::size_t j) const
	{
		return children_(i, j);
	}

	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::add_row()
	{
		set_row_count(row_count() + 1);
	}

	void grid_layout::add_row(std::shared_ptr<element> c)
	{
		set_row_count(row_count() + 1);
		set(row_count() - 1, 0, c);
	}

	void grid_layout::add_column()
	{
		set_column_count(column_count() + 1);
	}

	void grid_layout::add_column(std::shared_ptr<element> c)
	{
		set_column_count(column_count() + 1);
		set(0, column_count() - 1, c);
	}

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

	template <int Dimension>
	static std::vector<math::interval<float>> cell_sizes(util::ndarray<std::shared_ptr<element>, 2> const & children)
	{
		std::vector<math::interval<float>> result(children.dim(Dimension), {0.f, inf});

		for (auto idx : children.indices())
		{
			auto c = children(idx);
			if (!c) continue;

			result[idx[Dimension]] &= c->size_constraints()[1 - Dimension];
		}

		return result;
	}

	template <int Dimension>
	static std::vector<math::interval<float>> cell_sizes_constrained(util::ndarray<std::shared_ptr<element>, 2> const & children, std::vector<math::interval<float>> const & other_dimension_shapes)
	{
		std::vector<math::interval<float>> result(children.dim(Dimension), {0.f, inf});

		for (auto idx : children.indices())
		{
			auto c = children(idx);
			if (!c) continue;

			math::interval<float> constraint;

			if constexpr (Dimension == 0)
			{
				constraint = c->height_constraints(other_dimension_shapes[idx[1]].length());
			}
			else if constexpr (Dimension == 1)
			{
				constraint = c->width_constraints(other_dimension_shapes[idx[0]].length());
			}

			result[idx[Dimension]] &= constraint;
		}

		return result;
	}

	static math::interval<float> combine(std::vector<math::interval<float>> const & sizes, std::vector<float> const & weights, float margin, bool outer)
	{
		math::interval<float> result{0.f, 0.f};
		math::interval<float> unit{0.f, inf};

		float sum_weights = 0.f;

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

			if (minimized)
				result += sizes[i].min;
			else
				unit &= math::interval<float>{sizes[i].min / weights[i], sizes[i].max / weights[i]};

			sum_weights += weights[i];
		}

		// prevent (inf * 0.f)
		if (sum_weights > 0.f)
		{
			result.min += unit.min * sum_weights;
			result.max += unit.max * sum_weights;
		}

		int const margin_count = (std::max<int>(1, sizes.size()) + (outer ? 1 : -1));

		result += margin_count * margin;

		return result;
	}

	static std::vector<math::interval<float>> allocate(std::vector<math::interval<float>> const & sizes, std::vector<float> const & weights, float available, float margin, bool outer)
	{
		available -= margin * (static_cast<int>(sizes.size()) + (outer ? 1 : -1));

		float const sum_weights = std::accumulate(weights.begin(), weights.end(), 0.f);

		std::vector<math::interval<float>> result(sizes.size());

		// Allocate minimized
		for (std::size_t i = 0; i < sizes.size(); ++i)
		{
			if (weights[i] == 0.f)
			{
				result[i] = {0.f, sizes[i].min};
				available -= sizes[i].min;
			}
		}

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

		// Allocate scalable
		for (std::size_t i = 0; i < sizes.size(); ++i)
		{
			if (weights[i] == 0.f)
				continue;

			float const column_width = available * weights[i] / sum_weights;
			result[i] = {0.f, column_width};
		}

		return result;
	}

	static void move(std::vector<math::interval<float>> & shapes, math::interval<float> const & available, float margin, bool outer)
	{
		// Move cells one after another
		for (std::size_t i = 0; i < shapes.size(); ++i)
		{
			if (i == 0)
				shapes[i] += available.min + (outer ? margin : 0.f);
			else
				shapes[i] += shapes[i - 1].max + margin;
		}
	}

	math::box<float, 2> grid_layout::size_constraints() const
	{
		auto st = merged_own_style();

		auto row_sizes = cell_sizes<0>(children_);
		auto column_sizes = cell_sizes<1>(children_);

		auto height = combine(row_sizes, row_weight_, *st->outer_margin, outer_margin_);
		auto width = combine(column_sizes, column_weight_, *st->outer_margin, outer_margin_);
		return {{width, height}};
	}

	math::interval<float> grid_layout::width_constraints(float height) const
	{
		auto st = merged_own_style();

		auto row_sizes = cell_sizes<0>(children_);
		auto row_shapes = allocate(row_sizes, row_weight_, height, *st->outer_margin, outer_margin_);

		auto column_sizes = cell_sizes_constrained<1>(children_, row_shapes);

		return combine(column_sizes, column_weight_, *st->outer_margin, outer_margin_);
	}

	math::interval<float> grid_layout::height_constraints(float width) const
	{
		auto st = merged_own_style();

		auto column_sizes = cell_sizes<1>(children_);
		auto column_shapes = allocate(column_sizes, column_weight_, width, *st->outer_margin, outer_margin_);

		auto row_sizes = cell_sizes_constrained<0>(children_, column_shapes);

		return combine(row_sizes, row_weight_, *st->outer_margin, outer_margin_);
	}

	void grid_layout::reshape(math::box<float, 2> const & bbox)
	{
		shape_.box = bbox;

		auto st = merged_own_style();

		if (width_first_)
		{
			auto column_sizes = cell_sizes<1>(children_);

			column_shape_ = allocate(column_sizes, column_weight_, bbox[0].length(), *st->outer_margin, outer_margin_);

			auto row_sizes = cell_sizes_constrained<0>(children_, column_shape_);

			row_shape_ = allocate(row_sizes, row_weight_, bbox[1].length(), *st->outer_margin, outer_margin_);
		}
		else
		{
			auto row_sizes = cell_sizes<0>(children_);

			row_shape_ = allocate(row_sizes, row_weight_, bbox[1].length(), *st->outer_margin, outer_margin_);

			auto column_sizes = cell_sizes_constrained<1>(children_, row_shape_);

			column_shape_ = allocate(column_sizes, column_weight_, bbox[0].length(), *st->outer_margin, outer_margin_);
		}

		move(row_shape_, bbox[1], *st->outer_margin, outer_margin_);
		move(column_shape_, bbox[0], *st->outer_margin, outer_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();
	}

}