#pragma once

#include <psemek/ml/neural_net/neural_net.hpp>

namespace psemek::ml
{

	// A class that stores temporary data to facilitate
	//   allocation-free multiple evaluations of a neural net
	template <typename T>
	struct neural_net_evaluator
	{
		std::vector<T> apply(neural_net<T> const & nn, std::vector<T> input) const;

	private:
		mutable std::vector<T> temp_;
	};

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

	template <typename T>
	std::vector<T> neural_net_evaluator<T>::apply(neural_net<T> const & nn, std::vector<T> input) const
	{
		auto layer_sizes = nn.layer_sizes();
		auto weights = nn.weights().begin();

		if (layer_sizes[0] != input.size())
			throw wrong_neural_net_input_size{layer_sizes[0], input.size()};

		for (std::size_t l = 0; l + 1 < layer_sizes.size(); ++l)
		{
			temp_.resize(layer_sizes[l + 1]);

			for (std::size_t i = 0; i < layer_sizes[l + 1]; ++i)
			{
				temp_[i] = *weights++;

				for (std::size_t j = 0; j < layer_sizes[l]; ++j)
					temp_[i] += (*weights++) * input[j];

				temp_[i] = activation(temp_[i], nn.activation_types()[l]);
			}

			std::swap(temp_, input);
		}

		return input;
	}

}
