#pragma once

#include <psemek/util/function.hpp>
#include <psemek/util/ndarray.hpp>
#include <psemek/util/hash_table.hpp>
#include <psemek/math/box.hpp>

#include <optional>

namespace psemek::pcg
{

	template <typename T, std::size_t N>
	struct chunked_map
	{
		using generator_func = util::function<util::ndarray<T, N>(math::box<int, N> const &)>;

		chunked_map(std::size_t chunk_size, generator_func generator)
			: chunk_size_(chunk_size)
			, generator_(std::move(generator))
		{}

		T & operator()(math::point<int, N> p)
		{
			math::vector<int, N> ip;
			for (std::size_t i = 0; i < N; ++i)
			{
				ip[i] = (p[i] >= 0) ? (p[i] / chunk_size_) : -((- p[i] + chunk_size_ - 1) / chunk_size_);
			}

			auto it = chunks_.find(ip);
			if (it == chunks_.end())
			{
				math::box<int, 2> box;
				for (std::size_t i = 0; i < N; ++i)
				{
					box[i] = {ip[i] * chunk_size_, (ip[i] + 1) * chunk_size_};
				}
				it = chunks_.insert(std::pair{ip, generator_(box)}).first;
			}

			std::array<std::size_t, N> t;
			for (std::size_t i = 0; i < N; ++i)
				t[i] = p[i] - ip[i] * chunk_size_;

			return it->second(t);
		}

		std::optional<T> at(math::point<int, N> p) const
		{
			math::vector<int, N> ip;
			for (std::size_t i = 0; i < N; ++i)
			{
				ip[i] = (p[i] >= 0) ? (p[i] / chunk_size_) : -((- p[i] + chunk_size_ - 1) / chunk_size_);
			}

			auto it = chunks_.find(ip);
			if (it == chunks_.end())
				return std::nullopt;

			std::array<std::size_t, N> t;
			for (std::size_t i = 0; i < N; ++i)
				t[i] = p[i] - ip[i] * chunk_size_;

			return it->second(t);
		}

	private:
		int const chunk_size_;
		generator_func generator_;
		mutable util::hash_map<math::vector<int, N>, util::ndarray<T, N>> chunks_;

	};

}
