#include <psemek/test/test.hpp>

#include <psemek/ml/neural_net/learner.hpp>
#include <psemek/ml/neural_net/evaluator.hpp>
#include <psemek/ml/neural_net/randomize.hpp>
#include <psemek/ml/neural_net/loss.hpp>
#include <psemek/random/generator.hpp>
#include <psemek/random/uniform.hpp>
#include <psemek/geom/math.hpp>

using namespace psemek::ml;
using namespace psemek::random;
using namespace psemek::geom;

test_case(ml_neural__net_gradient)
{
	generator rng;
	for (std::size_t iteration = 0; iteration < 64; ++iteration)
	{
		std::vector<std::size_t> sizes;
		sizes.resize(uniform<std::size_t>(rng, 2, 5));
		for (auto & s : sizes)
			s = uniform<std::size_t>(rng, 1, 50);

		std::vector<activation_type> activations(sizes.size() - 1);
		for (auto & a : activations)
			a = static_cast<activation_type>(uniform<std::size_t>(rng, 0, static_cast<std::size_t>(activation_type_values().size()) - 1));

		neural_net<double> nn(std::move(sizes), std::move(activations));
		randomize_normal(nn, rng);

		std::vector<double> input(nn.layer_sizes().front());
		for (auto & x : input)
			x = uniform<double>(rng);

		std::vector<double> output(nn.layer_sizes().back());
		for (auto & x : output)
			x = uniform<double>(rng);

		neural_net_learner<double> learner;
		learner.apply(nn, input);
		learner.backpropagate_l2(nn, output);

		double const eps = 1e-6;

		neural_net_evaluator<double> evaluator;
		for (std::size_t i = 0; i < nn.weights().size(); ++i)
		{
			double old = nn.weights()[i];
			nn.weights()[i] -= eps;
			double v0 = l2_loss(evaluator.apply(nn, input), output);
			nn.weights()[i] += 2.0 * eps;
			double v1 = l2_loss(evaluator.apply(nn, input), output);
			nn.weights()[i] = old;

			double numeric_gradient = (v1 - v0) / 2.0 / eps;
			expect_close(numeric_gradient, learner.gradient()[i], 1e-4);
		}
	}
}

test_case(ml_neural__net_arg__gradient)
{
	generator rng;
	for (std::size_t iteration = 0; iteration < 64; ++iteration)
	{
		std::vector<std::size_t> sizes;
		sizes.resize(uniform<std::size_t>(rng, 2, 5));
		for (auto & s : sizes)
			s = uniform<std::size_t>(rng, 1, 50);

		std::vector<activation_type> activations(sizes.size() - 1);
		for (auto & a : activations)
			a = static_cast<activation_type>(uniform<std::size_t>(rng, 0, static_cast<std::size_t>(activation_type_values().size()) - 1));

		neural_net<double> nn(std::move(sizes), std::move(activations));
		randomize_normal(nn, rng);

		std::vector<double> input(nn.layer_sizes().front());
		for (auto & x : input)
			x = uniform<double>(rng);

		std::vector<double> output(nn.layer_sizes().back());
		for (auto & x : output)
			x = uniform<double>(rng);

		neural_net_learner<double> learner;
		learner.apply(nn, input);
		learner.backpropagate_l2(nn, output);

		double const eps = 1e-6;

		neural_net_evaluator<double> evaluator;
		for (std::size_t i = 0; i < input.size(); ++i)
		{
			double old = input[i];
			input[i] -= eps;
			double v0 = l2_loss(evaluator.apply(nn, input), output);
			input[i] += 2.0 * eps;
			double v1 = l2_loss(evaluator.apply(nn, input), output);
			input[i] = old;

			double numeric_gradient = (v1 - v0) / 2.0 / eps;
			expect_close(numeric_gradient, learner.arg_gradient()[i], 1e-4);
		}
	}
}