#pragma once

#include <psemek/util/ndarray.hpp>

#include <map>
#include <optional>

namespace psemek::util
{

	template <typename T, std::size_t N, typename Key, typename Compare = std::less<Key>>
	struct atlas;

	template <typename T, typename Key, typename Compare>
	struct atlas<T, 2, Key, Compare>
	{
		struct atlas_part
		{
			std::array<std::size_t, 2> begin;
			std::array<std::size_t, 2> end;
		};

		enum class padding_mode
		{
			default_value,
			clamp,
		};

		atlas(T default_value = T{}, std::size_t padding = 0, padding_mode mode = padding_mode::default_value, Compare compare = Compare{})
			: default_value_(std::move(default_value))
			, padding_(padding)
			, padding_mode_(mode)
			, data_(std::move(compare))
		{}

		util::ndarray<T, 2> const & array() const
		{
			return array_;
		}

		std::pair<atlas_part, bool> insert(Key const & key, util::ndarray<T, 2> const & data)
		{
			auto it = data_.find(key);
			if (it != data_.end())
				return {it->second, false};

			std::size_t part_width = data.width() + padding_ * 2;
			std::size_t part_height = data.height() + padding_ * 2;

			atlas_part part;
			part.begin[0] = free_start_;
			part.end[0] = part.begin[0] + part_width;
			part.begin[1] = 0;
			part.end[1] = part_height;

			std::size_t new_width = array_.width();
			std::size_t new_height = array_.height();

			while (new_width < free_start_ + part_width)
			{
				if (new_width == 0)
					new_width = 1;
				else
					new_width *= 2;
			}

			while (new_height < part_height)
			{
				if (new_height == 0)
					new_height = 1;
				else
					new_height *= 2;
			}

			array_.resize({new_width, new_height}, default_value_);

			free_start_ += part_width;

			for (std::size_t y = 0; y < part_height; ++y)
			{
				for (std::size_t x = 0; x < part_width; ++x)
				{
					T value;
					if (padding_mode_ == padding_mode::default_value)
					{
						if (x < padding_ || y < padding_ || x + padding_ >= part_width || y + padding_ >= part_height)
							value = default_value_;
						else
							value = data(x - padding_, y - padding_);
					}
					else
					{
						std::size_t xx = std::min(std::max(x, padding_), part_width - padding_ - 1);
						std::size_t yy = std::min(std::max(y, padding_), part_height - padding_ - 1);
						value = data(xx - padding_, yy - padding_);
					}

					array_(part.begin[0] + x, part.begin[1] + y) = value;
				}
			}

			part.begin[0] += padding_;
			part.begin[1] += padding_;
			part.end[0] -= padding_;
			part.end[1] -= padding_;

			data_[key] = part;
			return {part, true};
		}

		std::optional<atlas_part> find(Key const & key) const
		{
			auto it = data_.find(key);
			if (it != data_.end())
				return it->second;
			return std::nullopt;
		}

		atlas_part at(Key const & key) const
		{
			return data_.at(key);
		}

	private:
		T default_value_;
		std::size_t const padding_;
		padding_mode const padding_mode_;
		std::map<Key, atlas_part, Compare> data_;
		util::ndarray<T, 2> array_;
		std::size_t free_start_ = 0;
	};

}
