#include <psemek/ui/impl/box_layout_base.hpp>
#include <psemek/ui/impl/size_polygon_utils.hpp>
#include <psemek/react/map.hpp>
#include <psemek/react/join.hpp>

namespace psemek::ui::impl
{

	namespace
	{

		static constexpr box_layout::size_policy default_policy = box_layout::weight{};

		template <int Dimension>
		size_polygon compute_size_constraints(std::vector<box_layout::size_policy> const & size_policies,
			std::vector<size_polygon const *> const & children_size_constraints, float margin)
		{
			size_polygon minimized{{0.f, 0.f}, {0.f, max_size}};
			size_polygon weight_unit = minimized;
			float weight_sum = 0.f;

			for (std::size_t i = 0; i < size_policies.size(); ++i)
			{
				if (std::get_if<box_layout::minimized>(&size_policies[i]))
				{
					minimized = sum_along(minimized, *children_size_constraints[i], Dimension);
				}
				else if (auto weight = std::get_if<box_layout::weight>(&size_policies[i]))
				{
					size_polygon child = *children_size_constraints[i];
					for (auto & p : child)
						p[Dimension] /= weight->value;
					weight_unit = sum_along(weight_unit, child, Dimension);
					weight_sum += weight->value;
				}
			}

			auto result = std::move(minimized);

			if (weight_sum > 0.f)
			{
				for (auto & p : weight_unit)
					p[Dimension] *= weight_sum;
				result = sum_along(result, weight_unit, Dimension);
			}

			geom::vector shift_delta{0.f, 0.f};
			shift_delta[Dimension] = margin;

			result = shift(std::move(result), shift_delta);

			return result;
		}

		template <int Dimension>
		std::vector<float> allocate(float total_size, std::vector<react::value<box_layout::size_policy>> const & policies,
			util::span<std::unique_ptr<component> const> children)
		{
			std::vector<float> result(policies.size(), 0.f);

			float total_weight = 0.f;

			// First, allocate minimized elements
			// Also compute the total weight for weighted elements
			for (std::size_t i = 0; i < policies.size(); ++i)
			{
				auto const policy = policies[i] ? *policies[i] : default_policy;

				if (std::get_if<box_layout::minimized>(&policy))
				{
					result[i] = min(*children[i]->size_constraints(), Dimension).front()[Dimension];
					total_size -= result[i];
				}
				else if (auto weight = std::get_if<box_layout::weight>(&policy))
				{
					total_weight += weight->value;
				}
			}

			// Next, allocate weighted elements
			for (std::size_t i = 0; i < policies.size(); ++i)
			{
				auto const policy = policies[i] ? *policies[i] : default_policy;

				if (auto weight = std::get_if<box_layout::weight>(&policy))
				{
					result[i] = total_size * weight->value / total_weight;
				}
			}

			return result;
		}

	}

	template <int Dimension>
	box_layout_base<Dimension>::box_layout_base(react::value<float> margin)
		: margin_(margin)
		, size_policies_({})
		, children_size_constraints_()
		, size_constraints_(react::map([](auto const & size_policies, auto const & children_size_constraints, auto const & margin){
			return compute_size_constraints<Dimension>(size_policies, children_size_constraints, margin);
		},
			react::join(react::map(react::unpack_with_default(default_policy), size_policies_)),
			react::join(react::map(react::unpack_with_transform([](react::value<size_polygon> const & arg){ return &(*arg); }), children_size_constraints_)),
			margin_))
	{}

	template <int Dimension>
	void box_layout_base<Dimension>::reshape(geom::box<float, 2> const & new_shape)
	{
		container::reshape(new_shape);

		if (children().empty())
			return;

		float const margin = *margin_;

		auto sizes = allocate<Dimension>(new_shape[Dimension].length() - (children().size() - 1) * margin, *size_policies_, children());

		float pen = new_shape[Dimension].min;
		for (std::size_t i = 0; i < children().size(); ++i)
		{
			if (!children()[i]) continue;

			geom::box<float, 2> child_shape;
			child_shape[Dimension] = {pen, pen + sizes[i]};
			child_shape[Dimension ^ 1] = new_shape[Dimension ^ 1];
			children()[i]->reshape(child_shape);

			pen += sizes[i];
			pen += margin;
		}
	}

	template <int Dimension>
	react::value<size_polygon> box_layout_base<Dimension>::size_constraints() const
	{
		return size_constraints_;
	}

	template <int Dimension>
	void box_layout_base<Dimension>::set_children(std::vector<std::unique_ptr<component>> children)
	{
		container::set_children(std::move(children));

		std::vector<react::value<size_polygon>> children_size_constraints;
		for (auto const & child : this->children())
		{
			if (child)
				children_size_constraints.push_back(child->size_constraints());
			else
				children_size_constraints.push_back(default_size_constraints());
		}
		children_size_constraints_.set(std::move(children_size_constraints));
	}

	template <int Dimension>
	std::vector<std::unique_ptr<component>> box_layout_base<Dimension>::release_children()
	{
		auto result = container::release_children();
		children_size_constraints_.set({});
		return result;
	}

	template <int Dimension>
	void box_layout_base<Dimension>::update(typename detail::box_layout_type<Dimension>::type const & box_layout)
	{
		size_policies_ = react::map([](auto const & children){
			std::vector<react::value<box_layout::size_policy>> size_policies;
			size_policies.reserve(children.size());
			for (auto const & child : children)
				size_policies.push_back(child.policy);
			return size_policies;
		}, box_layout.children);
	}

	template struct box_layout_base<0>;
	template struct box_layout_base<1>;

}