#pragma once

#include <psemek/util/ndarray.hpp>
#include <psemek/util/exception.hpp>
#include <psemek/util/to_string.hpp>

#include <cstdint>

namespace psemek::util
{

	template <typename T, std::size_t N, typename Index = std::int32_t>
	struct spatial_array
	{
		spatial_array();

		using index_type = Index;

		template <typename ... Ix>
		T & get(Ix ... index);

		template <typename ... Ix>
		T * get_if(Ix ... index);

		template <typename ... Ix>
		T const * get_if(Ix ... index) const;

		template <typename ... Ix>
		T const & at(Ix ... index) const;

		template <typename ... Ix>
		bool contains(Ix ... index) const;

		template <typename ... Ix>
		void expand(Ix ... index);

		void clear();
		bool empty() const;

		T * begin();
		T * end();

		T const * begin() const;
		T const * end() const;

		Index size(std::size_t dimension) const;
		Index min(std::size_t dimension) const;
		Index max(std::size_t dimension) const;

	private:
		Index origin_[N];
		ndarray<T, N> array_;
	};

	template <typename T, std::size_t N, typename Index>
	spatial_array<T, N, Index>::spatial_array()
	{
		for (std::size_t i = 0; i < N; ++i)
			origin_[i] = 0;
	}

	template <typename T, std::size_t N, typename Index>
	template <typename ... Ix>
	T & spatial_array<T, N, Index>::get(Ix ... index)
	{
		static_assert(sizeof...(Ix) == N, "wrong number of indices");

		if (!contains(index...))
			expand(index...);

		Index ix[N] {index...};
		for (std::size_t i = 0; i < N; ++i)
			ix[i] -= origin_[i];
		return array_(ix);
	}

	template <typename T, std::size_t N, typename Index>
	template <typename ... Ix>
	T * spatial_array<T, N, Index>::get_if(Ix ... index)
	{
		static_assert(sizeof...(Ix) == N, "wrong number of indices");

		if (!contains(index...))
			return nullptr;

		Index ix[N] {index...};
		for (std::size_t i = 0; i < N; ++i)
			ix[i] -= origin_[i];
		return &array_(ix);
	}

	template <typename T, std::size_t N, typename Index>
	template <typename ... Ix>
	T const * spatial_array<T, N, Index>::get_if(Ix ... index) const
	{
		static_assert(sizeof...(Ix) == N, "wrong number of indices");

		if (!contains(index...))
			return nullptr;

		Index ix[N] {index...};
		for (std::size_t i = 0; i < N; ++i)
			ix[i] -= origin_[i];
		return &array_(ix);
	}

	template <typename T, std::size_t N, typename Index>
	template <typename ... Ix>
	T const & spatial_array<T, N, Index>::at(Ix ... index) const
	{
		static_assert(sizeof...(Ix) == N, "wrong number of indices");

		auto const size = array_.dims();
		Index ix[N] {index...};
		for (std::size_t i = 0; i < N; ++i)
		{
			if (ix[i] < origin_[i] || ix[i] >= origin_[i] + static_cast<Index>(size[i]))
				throw exception(to_string("Spatial array index ", ix[i], "(#", i, ") is out of bounds [", origin_[i], origin_[i] + static_cast<Index>(size[i]), ")"));
			ix[i] -= origin_[i];
		}
		return array_(ix);
	}

	template <typename T, std::size_t N, typename Index>
	template <typename ... Ix>
	bool spatial_array<T, N, Index>::contains(Ix ... index) const
	{
		static_assert(sizeof...(Ix) == N, "wrong number of indices");

		auto const size = array_.dims();
		Index ix[N] {index...};
		for (std::size_t i = 0; i < N; ++i)
			if (ix[i] < origin_[i] || ix[i] >= origin_[i] + static_cast<Index>(size[i]))
				return false;
		return true;
	}

	template <typename T, std::size_t N, typename Index>
	template <typename ... Ix>
	void spatial_array<T, N, Index>::expand(Ix ... index)
	{
		static_assert(sizeof...(Ix) == N, "wrong number of indices");

		Index ix[N] {index...};

		if (array_.empty())
		{
			std::array<std::size_t, N> size;

			for (std::size_t i = 0; i < N; ++i)
			{
				origin_[i] = ix[i];
				size[i] = 1;
			}

			array_.resize(size);
			*array_.begin() = T();
			return;
		}

		Index new_origin[N];
		std::array<std::size_t, N> new_size = array_.dims();

		for (std::size_t i = 0; i < N; ++i)
			new_origin[i] = origin_[i];

		for (std::size_t i = 0; i < N; ++i)
		{
			while (ix[i] >= new_origin[i] + static_cast<Index>(new_size[i]))
				new_size[i] += new_size[i];

			while (ix[i] < new_origin[i])
			{
				new_origin[i] -= static_cast<Index>(new_size[i]);
				new_size[i] += new_size[i];
			}
		}

		ndarray<T, N> new_array(new_size);
		for (auto & v : new_array)
			v = T();

		for (auto idx : array_.indices())
		{
			std::array<std::size_t, N> new_idx;
			for (std::size_t i = 0; i < N; ++i)
				new_idx[i] = static_cast<Index>(idx[i]) + origin_[i] - new_origin[i];
			new_array(new_idx) = std::move(array_(idx));
		}

		array_ = std::move(new_array);

		for (std::size_t i = 0; i < N; ++i)
			origin_[i] = new_origin[i];
	}

	template <typename T, std::size_t N, typename Index>
	void spatial_array<T, N, Index>::clear()
	{
		array_.clear();
	}

	template <typename T, std::size_t N, typename Index>
	bool spatial_array<T, N, Index>::empty() const
	{
		return array_.empty();
	}

	template <typename T, std::size_t N, typename Index>
	T * spatial_array<T, N, Index>::begin()
	{
		return array_.begin();
	}

	template <typename T, std::size_t N, typename Index>
	T * spatial_array<T, N, Index>::end()
	{
		return array_.end();
	}

	template <typename T, std::size_t N, typename Index>
	T const * spatial_array<T, N, Index>::begin() const
	{
		return array_.begin();
	}

	template <typename T, std::size_t N, typename Index>
	T const * spatial_array<T, N, Index>::end() const
	{
		return array_.end();
	}

	template <typename T, std::size_t N, typename Index>
	Index spatial_array<T, N, Index>::size(std::size_t dimension) const
	{
		return static_cast<Index>(array_.dim(dimension));
	}

	template <typename T, std::size_t N, typename Index>
	Index spatial_array<T, N, Index>::min(std::size_t dimension) const
	{
		return origin_[dimension];
	}

	template <typename T, std::size_t N, typename Index>
	Index spatial_array<T, N, Index>::max(std::size_t dimension) const
	{
		return origin_[dimension] + static_cast<Index>(array_.dim(dimension));
	}

}
