#pragma once

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

#include <psemek/math/vector.hpp>
#include <psemek/math/box.hpp>

#include <optional>

namespace psemek::pcg
{

	namespace detail
	{

		template <typename T, std::size_t N, typename F, typename G>
		T perlin_impl(math::point<T, N> p, int grid_size, F && grid_at, G && bound)
		{
			math::vector<int, N> ip;
			for (std::size_t i = 0; i < N; ++i)
			{
				ip[i] = (p[i] >= 0) ? (p[i] / grid_size) : -((- p[i] + grid_size - 1) / grid_size);
			}
			ip = bound(ip);

			math::vector<T, N> t = (p - p.zero() - math::cast<T>(ip) * T(grid_size)) / T(grid_size);

			T values[1 << N];

			for (std::size_t mask = 0; mask < (1 << N); ++mask)
			{
				math::vector<T, N> tt;
				math::vector<int, N> ii;
				for (std::size_t i = 0; i < N; ++i)
				{
					bool m = (mask & (1 << i));
					tt[i] = m ? t[i] - 1 : t[i];
					ii[i] = m ? ip[i] + 1 : ip[i];
				}
				values[mask] = math::dot(grid_at(ii), tt);
			}

			auto smootherstep = [](T x0, T x1, T t)
			{
				auto const s = t * t * t * (T(10) + t * (-T(15) + T(6) * t));
				return x0 * (T(1) - s) + x1 * s;
			};

			for (std::size_t i = N; i --> 0;)
			{
				for (std::size_t mask = 0; mask < (1 << i); ++mask)
					values[mask] = smootherstep(values[mask], values[mask | (1 << i)], t[i]);
			}

			static auto const max = std::sqrt(T(N) / T(4));
			return 0.5 * (1. + values[0] / max);
		}

	}

	template <typename T, std::size_t N>
	struct lazy_perlin_view
	{
		using value_type = T;
		static constexpr auto dimension = N;

		lazy_perlin_view() = default;

		lazy_perlin_view(std::size_t grid_size, math::vector<int, N> const & view_origin, util::ndarray<math::vector<T, N>, N> subgrid)
			: grid_size_(grid_size)
			, origin_(view_origin)
			, subgrid_(std::move(subgrid))
		{}

		T operator()(math::point<T, N> p) const;

	private:
		int grid_size_;
		math::vector<int, N> origin_;
		util::ndarray<math::vector<T, N>, N> subgrid_;
	};

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

		using value_type = T;
		static constexpr auto dimension = N;

		lazy_perlin(std::size_t grid_size, generator_func generator)
			: grid_size_(grid_size)
			, gen_(std::move(generator))
		{}

		T operator()(math::point<T, N> p) const;

		auto subview(math::box<T, N> const & box) const;

	private:
		int const grid_size_;
		mutable util::hash_map<math::vector<int, N>, math::vector<T, N>> grid_;
		generator_func gen_;

		math::vector<T, N> grid_at(math::vector<int, N> const & c) const
		{
			auto it = grid_.find(c);
			if (it == grid_.end())
				it = grid_.insert({c, gen_(c)}).first;
			return it->second;
		}
	};

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

		using value_type = T;
		static constexpr auto dimension = N;

		bounded_lazy_perlin(std::size_t grid_size, generator_func generator, math::box<int, N> const & bounds)
			: grid_size_(grid_size)
			, gen_(std::move(generator))
		{
			std::array<std::size_t, N> dims;

			for (std::size_t i = 0; i < N; ++i)
			{
				origin_[i] = bounds[i].min / grid_size;
				dims[i] = bounds[i].length() / grid_size + 1;
				corner_[i] = origin_[i] + dims[i] - 2;
			}

			grid_.resize(dims);
		}

		T operator()(math::point<T, N> p) const;

	private:
		int const grid_size_;
		math::vector<int, N> origin_;
		math::vector<int, N> corner_;
		mutable util::ndarray<std::optional<math::vector<T, N>>, N> grid_;
		generator_func gen_;

		math::vector<T, N> grid_at(math::vector<int, N> const & c) const
		{
			std::array<std::size_t, N> idx;
			for (std::size_t i = 0; i < N; ++i)
				idx[i] = c[i] - origin_[i];
			auto & result = grid_(idx);
			if (!result)
				result = gen_(c);
			return *result;
		}
	};

	template <typename T, std::size_t N>
	T lazy_perlin_view<T, N>::operator()(math::point<T, N> p) const
	{
		return detail::perlin_impl(p, grid_size_, [this](auto const & c){
			return subgrid_(c - origin_);
		}, [this](auto p){
			for (std::size_t i = 0; i < N; ++i)
			{
				p[i] = std::max<int>(p[i], origin_[i]);
				p[i] = std::min<int>(p[i], origin_[i] + subgrid_.dim(i));
			}
			return p;
		});
	}

	template <typename T, std::size_t N>
	T lazy_perlin<T, N>::operator()(math::point<T, N> p) const
	{
		return detail::perlin_impl(p, grid_size_, [this](auto const & c){ return grid_at(c); }, [](auto const & p){ return p; });
	}

	template <typename T, std::size_t N>
	auto lazy_perlin<T, N>::subview(math::box<T, N> const & box) const
	{
		math::vector<int, N> subgrid_origin;
		std::array<std::size_t, N> subgrid_dimensions;
		for (std::size_t i = 0; i < N; ++i)
		{
			subgrid_origin[i] = std::floor(box[i].min / grid_size_);
			int subgrid_max = std::ceil(box[i].max / grid_size_);
			subgrid_dimensions[i] = subgrid_max - subgrid_origin[i] + 1;
		}

		util::ndarray<math::vector<T, N>, N> subgrid(subgrid_dimensions);
		for (auto const idx : subgrid.indices())
		{
			math::vector<int, N> id;
			for (std::size_t i = 0; i < N; ++i)
				id[i] = subgrid_origin[i] + idx[i];
			subgrid(idx) = grid_at(id);
		}

		return lazy_perlin_view<T, N>(grid_size_, subgrid_origin, std::move(subgrid));
	}

	template <typename T, std::size_t N>
	T bounded_lazy_perlin<T, N>::operator()(math::point<T, N> p) const
	{
		return detail::perlin_impl(p, grid_size_, [this](auto const & c){ return grid_at(c); }, [this](auto p){
			for (std::size_t i = 0; i < N; ++i)
			{
				p[i] = std::max(p[i], origin_[i]);
				p[i] = std::min(p[i], corner_[i]);
			}
			return p;
		});
	}

}
