#pragma once

#include <psemek/ml/neural_net/activation.hpp>
#include <psemek/ml/neural_net/error.hpp>
#include <psemek/util/span.hpp>

#include <vector>

namespace psemek::ml
{

	namespace detail
	{

		inline std::pair<std::vector<std::size_t>, std::vector<activation_type>> make_nn_ctor_args(std::vector<std::size_t> layer_sizes, activation_type type)
		{
			if (layer_sizes.empty())
				throw empty_neural_net_error{};
			std::vector<activation_type> activation_types(layer_sizes.size() - 1, type);
			return {std::move(layer_sizes), std::move(activation_types)};
		}

		inline std::size_t weight_count(std::vector<std::size_t> const & layer_sizes)
		{
			std::size_t result = 0;
			for (std::size_t l = 0; l + 1 < layer_sizes.size(); ++l)
				result += (layer_sizes[l] + 1) * layer_sizes[l + 1];
			return result;
		}

		template <typename NN, typename Fn>
		void for_each_layer_impl(NN & nn, Fn && fn)
		{
			auto const layer_sizes = nn.layer_sizes();
			auto const weights = nn.weights();
			std::size_t offset = 0;
			for (std::size_t l = 0; l + 1 < layer_sizes.size(); ++l)
			{
				std::size_t count = (layer_sizes[l] + 1) * layer_sizes[l + 1];
				fn(l, util::span{weights.data() + offset, weights.data() + offset + count});
				offset += count;
			}
		}

	}

	template <typename T>
	struct neural_net
	{
		neural_net() = default;
		neural_net(std::vector<std::size_t> layer_sizes);
		neural_net(std::vector<std::size_t> layer_sizes, activation_type type);
		neural_net(std::vector<std::size_t> layer_sizes, std::vector<activation_type> activation_types);

		// A non-empty neural net is basically unusable
		bool empty() const { return layer_sizes_.empty(); }

		util::span<std::size_t const> layer_sizes() const { return layer_sizes_; }

		util::span<activation_type const> activation_types() const { return activation_types_; }

		// Weights are stored in a sequential manner, in the order in which they
		//   appear in evalution of the neural net, i.e. first come the weights
		//   between layers 0 and 1 (including the bias) in row-major order, then
		//   1-2 in row-major order, etc.
		util::span<T const> weights() const { return weights_; }
		util::span<T> weights() { return weights_; }

	private:
		std::vector<std::size_t> layer_sizes_;
		std::vector<activation_type> activation_types_;
		std::vector<T> weights_;

		// proxy constructor to overcome unspecified evaluation order in
		//   neural_net(std::move(layer_sizes), std::vector{layer_sizes.size(), activation_type::tanh})
		neural_net(std::pair<std::vector<std::size_t>, std::vector<activation_type>> args);

		void assert_nonempty() const
		{
			if (empty())
				throw empty_neural_net_error{};
		}
	};

	extern template struct neural_net<float>;
	extern template struct neural_net<double>;

	template <typename T>
	neural_net<T>::neural_net(std::vector<std::size_t> layer_sizes)
		: neural_net(std::move(layer_sizes), activation_type::tanh)
	{}

	template <typename T>
	neural_net<T>::neural_net(std::vector<std::size_t> layer_sizes, activation_type type)
		: neural_net(detail::make_nn_ctor_args(std::move(layer_sizes), type))
	{}

	template <typename T>
	neural_net<T>::neural_net(std::vector<std::size_t> layer_sizes, std::vector<activation_type> activation_types)
		: layer_sizes_(layer_sizes)
		, activation_types_(activation_types)
		, weights_(detail::weight_count(layer_sizes_))
	{
		if (layer_sizes_.empty())
			throw empty_neural_net_error{};

		if (layer_sizes_.size() != activation_types_.size() + 1)
			throw wrong_activation_types_count_error{layer_sizes_.size() - 1, activation_types_.size()};
	}

	template <typename T>
	neural_net<T>::neural_net(std::pair<std::vector<std::size_t>, std::vector<activation_type>> args)
		: neural_net(std::move(args.first), std::move(args.second))
	{}

	template <typename T, typename Fn>
	void for_each_layer(neural_net<T> const & nn, Fn && fn)
	{
		detail::for_each_layer_impl(nn, fn);
	}

	template <typename T, typename Fn>
	void for_each_layer(neural_net<T> & nn, Fn && fn)
	{
		detail::for_each_layer_impl(nn, fn);
	}

}
