#pragma once

#include <psemek/util/assert.hpp>

#include <memory>
#include <algorithm>
#include <array>
#include <type_traits>

namespace psemek::util
{

	namespace detail
	{

		template <typename T, typename = void>
		struct is_array : std::false_type {};

		template <typename T>
		struct is_array<T, std::void_t<decltype(std::declval<T&>()[0])>>
			: std::true_type
		{};

		template <typename T>
		constexpr bool is_array_v = is_array<T>::value;

		template <typename ... Ts>
		struct is_index;

		template <>
		struct is_index<>
			: std::false_type
		{};

		template <typename T, typename ... Ts>
		struct is_index<T, Ts...>
			: std::bool_constant<!is_array_v<T>>
		{};

		template <typename ... Ts>
		constexpr bool is_index_v = is_index<Ts...>::value;

	}

	template <typename T, std::size_t N>
	struct ndarray
	{
		static_assert(N >= 1);

		using value_type = T;
		static constexpr std::size_t dimension = N;

		using dims_type = std::array<std::size_t, N>;

		ndarray();
		ndarray(dims_type const & dims);
		ndarray(dims_type const & dims, T const & value);
		ndarray(dims_type const & dims, std::unique_ptr<T[]> data);
		ndarray(ndarray &&);

		ndarray(ndarray const &) = delete;

		ndarray & operator = (ndarray &&);

		ndarray & operator = (ndarray const &) = delete;

		dims_type const & dims() const
		{
			return dims_;
		}

		std::size_t dim(std::size_t i) const
		{
			assert(i < N);
			return dims_[i];
		}

		std::size_t size() const;

		std::size_t width() const
		{
			static_assert(N >= 1);
			return dims_[0];
		}

		std::size_t height() const
		{
			static_assert(N >= 2);
			return dims_[1];
		}

		std::size_t depth() const
		{
			static_assert(N >= 3);
			return dims_[2];
		}

		template <typename Index>
		std::enable_if_t<detail::is_array_v<Index>, T &>
			operator()(Index const & index);

		template <typename Index>
		std::enable_if_t<detail::is_array_v<Index>, T const &>
			operator()(Index const & index) const;

		template <typename ... Ixs>
		std::enable_if_t<detail::is_index_v<Ixs...>, T &>
			operator()(Ixs ... ixs);

		template <typename ... Ixs>
		std::enable_if_t<detail::is_index_v<Ixs...>, T const &>
			operator()(Ixs ... ixs) const;

		template <typename I>
		T & operator()(std::initializer_list<I> const & index);

		template <typename I>
		T const & operator()(std::initializer_list<I> const & index) const;

		ndarray copy() const;
		ndarray subarray(dims_type const & start, dims_type const & end) const;

		void resize(dims_type const & dims);

		void resize(dims_type const & dims, T const & value);

		void assign(dims_type const & dims, T const & value);

		T * data() { return data_.get(); }
		T const * data() const { return data_.get(); }

		T * begin() { return data_.get(); }
		T const * begin() const { return data_.get(); }
		T const * cbegin() const { return data_.get(); }

		T * end() { return data_.get() + size(); }
		T const * end() const { return data_.get() + size(); }
		T const * cend() const { return data_.get() + size(); }

		std::unique_ptr<T[]> release();

		bool empty() const;

		void clear();

		void fill(T const & value);

		auto index_begin() const;
		auto index_end() const;
		auto indices() const;

	private:
		std::unique_ptr<T[]> data_;
		std::array<std::size_t, N> dims_;

		void resize_impl(std::unique_ptr<T[]> data, dims_type const & dims);
	};

	namespace detail
	{

		template <std::size_t N>
		std::size_t product(std::array<std::size_t, N> const & dims)
		{
			std::size_t r = 1;
			for (auto const & d : dims)
				r *= d;
			return r;
		}

		template <std::size_t N>
		bool empty(std::array<std::size_t, N> const & dims)
		{
			for (auto const & d : dims)
				if (d == 0) return true;
			return false;
		}

		template <std::size_t N, typename Index>
		std::size_t index(Index const & i, std::array<std::size_t, N> const & dims)
		{
			std::size_t r = 0;
			for (std::size_t d = N; d --> 0;)
			{
				assert(i[d] < dims[d]);
				r = i[d] + r * dims[d];
			}
			return r;
		}

		template <std::size_t N>
		bool next(std::array<std::size_t, N> & i, std::array<std::size_t, N> const & dims)
		{
			for (std::size_t d = 0; d < N; ++d)
			{
				++i[d];
				if (i[d] < dims[d])
					return true;
				i[d] = 0;
			}
			return false;
		}

		template <std::size_t N>
		struct array_index_iterator
		{
			std::array<std::size_t, N> idx;
			std::array<std::size_t, N> dims;

			bool end;

			std::array<std::size_t, N> const & operator*() const { return idx; }

			array_index_iterator & operator ++()
			{
				end = !next(idx, dims);
				return *this;
			}
		};

		struct array_index_iterator_sentinel
		{};

		template <std::size_t N>
		bool operator == (array_index_iterator<N> const & it, array_index_iterator_sentinel)
		{
			return it.end;
		}

		template <std::size_t N>
		bool operator != (array_index_iterator<N> const & it, array_index_iterator_sentinel)
		{
			return !it.end;
		}

		template <std::size_t N>
		struct array_index_range
		{
			std::array<std::size_t, N> dims;

			array_index_iterator<N> begin() const
			{
				std::array<std::size_t, N> idx;
				for (std::size_t i = 0; i < N; ++i)
					idx[i] = 0;
				return {idx, dims, std::any_of(dims.begin(), dims.end(), [](std::size_t d){ return d == 0; })};
			}

			array_index_iterator_sentinel end() const
			{
				return {};
			}
		};

	}

	template <typename T, std::size_t N>
	auto ndarray<T, N>::index_begin() const
	{
		return indices().begin();
	}

	template <typename T, std::size_t N>
	auto ndarray<T, N>::index_end() const
	{
		return indices().end();
	}

	template <typename T, std::size_t N>
	auto ndarray<T, N>::indices() const
	{
		return detail::array_index_range<N>{dims_};
	}

	template <typename T, std::size_t N>
	ndarray<T, N>::ndarray()
	{
		dims_.fill(0);
	}

	template <typename T, std::size_t N>
	ndarray<T, N>::ndarray(dims_type const & dims)
		: dims_{dims}
	{
		data_.reset(new T[size()]);
	}

	template <typename T, std::size_t N>
	ndarray<T, N>::ndarray(dims_type const & dims, T const & value)
		: ndarray(dims)
	{
		fill(value);
	}

	template <typename T, std::size_t N>
	ndarray<T, N>::ndarray(dims_type const & dims, std::unique_ptr<T[]> data)
		: dims_{dims}
	{
		data_ = std::move(data);
	}

	template <typename T, std::size_t N>
	ndarray<T, N>::ndarray(ndarray && other)
		: data_{std::move(other.data_)}
		, dims_{other.dims_}
	{
		other.dims_.fill(0);
	}

	template <typename T, std::size_t N>
	ndarray<T, N> & ndarray<T, N>::operator = (ndarray && other)
	{
		if (this == &other)
			return *this;

		data_ = std::move(other.data_);
		dims_ = other.dims_;
		other.dims_.fill(0);

		return *this;
	}

	template <typename T, std::size_t N>
	std::size_t ndarray<T, N>::size() const
	{
		return detail::product(dims_);
	}

	template <typename T, std::size_t N>
	template <typename Index>
	std::enable_if_t<detail::is_array_v<Index>, T &>
		ndarray<T, N>::operator()(Index const & index)
	{
		return data_[detail::index(index, dims_)];
	}

	template <typename T, std::size_t N>
	template <typename Index>
	std::enable_if_t<detail::is_array_v<Index>, T const &>
		ndarray<T, N>::operator()(Index const & index) const
	{
		return data_[detail::index(index, dims_)];
	}

	template <typename T, std::size_t N>
	template <typename ... Ixs>
	std::enable_if_t<detail::is_index_v<Ixs...>, T &>
		ndarray<T, N>::operator()(Ixs ... ixs)
	{
		static_assert(sizeof...(Ixs) == N);
		dims_type dims{static_cast<std::size_t>(ixs)...};
		return (*this)(dims);
	}

	template <typename T, std::size_t N>
	template <typename ... Ixs>
	std::enable_if_t<detail::is_index_v<Ixs...>, T const &>
		ndarray<T, N>::operator()(Ixs ... ixs) const
	{
		static_assert(sizeof...(Ixs) == N);
		dims_type dims{static_cast<std::size_t>(ixs)...};
		return (*this)(dims);
	}

	template <typename T, std::size_t N>
	template <typename I>
	T & ndarray<T, N>::operator()(std::initializer_list<I> const & index)
	{
		return data_[detail::index(std::data(index), dims_)];
	}

	template <typename T, std::size_t N>
	template <typename I>
	T const & ndarray<T, N>::operator()(std::initializer_list<I> const & index) const
	{
		return data_[detail::index(std::data(index), dims_)];
	}

	template <typename T, std::size_t N>
	ndarray<T, N> ndarray<T, N>::copy() const
	{
		if (empty())
			return {};

		std::unique_ptr<T[]> data(new T[size()]);
		std::copy(begin(), end(), data.get());
		return ndarray{dims_, std::move(data)};
	}

	template <typename T, std::size_t N>
	ndarray<T, N> ndarray<T, N>::subarray(dims_type const & start, dims_type const & end) const
	{
		dims_type size;
		for (std::size_t i = 0; i < N; ++i)
			size[i] = end[i] - start[i];

		ndarray<T, N> result(size);
		for (auto const & idx : result.indices())
		{
			auto jdx = idx;
			for (std::size_t i = 0; i < N; ++i)
				jdx[i] += start[i];
			result(idx) = (*this)(jdx);
		}

		return result;
	}

	template <typename T, std::size_t N>
	void ndarray<T, N>::resize(dims_type const & dims)
	{
		if (dims == dims_)
			return;

		std::unique_ptr<T[]> data(new T[detail::product(dims)]);
		resize_impl(std::move(data), dims);
	}

	template <typename T, std::size_t N>
	void ndarray<T, N>::resize(dims_type const & dims, T const & value)
	{
		if (dims == dims_)
			return;

		auto const size = detail::product(dims);
		std::unique_ptr<T[]> data(new T[size]);
		std::fill(data.get(), data.get() + size, value);
		resize_impl(std::move(data), dims);
	}

	template <typename T, std::size_t N>
	void ndarray<T, N>::assign(dims_type const & dims, T const & value)
	{
		auto const size = detail::product(dims);
		std::unique_ptr<T[]> data(new T[size]);
		std::fill(data.get(), data.get() + size, value);
		data_ = std::move(data);
		dims_ = dims;
	}

	template <typename T, std::size_t N>
	std::unique_ptr<T[]> ndarray<T, N>::release()
	{
		dims_.fill(0);
		return std::move(data_);
	}

	template <typename T, std::size_t N>
	bool ndarray<T, N>::empty() const
	{
		return detail::empty(dims_);
	}

	template <typename T, std::size_t N>
	void ndarray<T, N>::clear()
	{
		data_.reset();
		dims_.fill(0);
	}

	template <typename T, std::size_t N>
	void ndarray<T, N>::fill(T const & value)
	{
		std::fill(data_.get(), data_.get() + size(), value);
	}

	template <typename T, std::size_t N>
	void ndarray<T, N>::resize_impl(std::unique_ptr<T[]> data, dims_type const & dims)
	{
		std::array<std::size_t, N> min_dim;
		for (std::size_t d = 0; d < N; ++d)
			min_dim[d] = std::min(dims_[d], dims[d]);
		if (!detail::empty(min_dim))
		{
			std::array<std::size_t, N> i;
			i.fill(0);

			do
			{
				data[detail::index(i, dims)] = std::move((*this)(i));
			} while (detail::next(i, min_dim));
		}

		dims_ = dims;
		data_ = std::move(data);
	}

	template <typename T, std::size_t N>
	void mirror(ndarray<T, N> & a, std::size_t i)
	{
		assert(i < N);

		auto const s = a.dim(i);
		for (auto idx : a.indices())
		{
			auto jdx = idx;
			jdx[i] = s - 1 - jdx[i];
			if (jdx[i] < idx[i])
				std::swap(a(idx), a(jdx));
		}
	}

	template <typename F, typename T, std::size_t N>
	auto map(F && f, ndarray<T, N> const & a)
	{
		using R = std::decay_t<decltype(f(T{}))>;

		ndarray<R, N> r(a.dims());

		auto begin = a.begin();
		auto end = a.end();
		auto out = r.begin();

		for (; begin != end;)
			*out++ = f(*begin++);

		return r;
	}

}
