#include <psemek/app/app.hpp>
#include <psemek/app/main.hpp>

#include <psemek/gfx/painter.hpp>

#include <psemek/geom/camera.hpp>
#include <psemek/geom/gauss.hpp>
#include <psemek/geom/distance.hpp>
#include <psemek/geom/interval.hpp>
#include <psemek/geom/math.hpp>
#include <psemek/geom/rotation.hpp>

#include <psemek/random/device.hpp>
#include <psemek/random/generator.hpp>
#include <psemek/random/uniform.hpp>
#include <psemek/random/normal.hpp>

#include <psemek/async/threadpool.hpp>

#include <psemek/util/clock.hpp>
#include <psemek/util/function.hpp>
#include <psemek/util/to_string.hpp>
#include <psemek/util/pretty_print.hpp>

//#include <eigen3/Eigen/Core>
//#include <eigen3/Eigen/Cholesky>

#include <fstream>
#include <iomanip>

using namespace psemek;

namespace
{

struct bone
{
	geom::point<float, 2> position;
	geom::vector<float, 2> direction;
	geom::vector<float, 2> velocity;
	float angular_velocity;
	float length;
	float width;
	float mass = 0.f;
	float inertia = 0.f;
};

struct joint
{
	std::size_t i0, i1;
	float s0, s1;
	// angle range max should be in [0, 2*pi]
	// angle range min should be less or equal to angle range max
	std::optional<geom::interval<float>> angle_range = std::nullopt;
};

template <std::size_t Bones, std::size_t Constraints>
struct constraint
{
	std::size_t bone[Bones];
	geom::matrix<float, Constraints, 4 * Bones> jacobian;
	geom::vector<float, Constraints> bias;
	geom::interval<float> range = geom::interval<float>::full();
};

struct system
{
	float const dt = 0.01f;

	std::vector<bone> bones;
	std::vector<joint> joints;

	struct selection
	{
		std::size_t index;
		float pos;
		geom::vector<float, 2> delta;
	};

	void advance(float time, std::optional<selection> sel, util::function<std::vector<float>(system const &)> torque_provider);

private:
	float time_lag_ = 0.f;
	std::vector<float> old_torque_;

	void step(std::optional<selection> sel, util::function<std::vector<float>(system const &)> const & torque_provider);

	template <std::size_t B, std::size_t C>
	void solve_constraint(constraint<B, C> const & c);
};

void system::advance(float time, std::optional<selection> sel, util::function<std::vector<float>(system const &)> torque_provider)
{
	time_lag_ += time;

	while (time_lag_ >= dt)
	{
		time_lag_ -= dt;
		step(sel, torque_provider);
	}
}

void system::step(std::optional<selection> sel, util::function<std::vector<float>(system const &)> const & torque_provider)
{
	if (old_torque_.size() != joints.size())
		old_torque_.assign(joints.size(), 0.f);

	geom::vector<float, 2> const gravity { 0.f, -10.f };

	for (std::size_t i = 0; i < bones.size(); ++i)
	{
		auto & b = bones[i];

		b.velocity += gravity * dt;
	}

	auto torque = torque_provider(*this);

	for (std::size_t i = 0; i < joints.size(); ++i)
	{
		auto const & j = joints[i];

		auto & b0 = bones[j.i0];
		auto & b1 = bones[j.i1];

		float const max_torque_change = 250.f * dt;

		float d = std::max(std::abs(torque[i] - old_torque_[i]) / max_torque_change, 1.f);

		old_torque_[i] = old_torque_[i] + (torque[i] - old_torque_[i]) / d;

		b0.angular_velocity -= old_torque_[i] * dt / b0.inertia;
		b1.angular_velocity += old_torque_[i] * dt / b1.inertia;
	}

	for (std::size_t i = 0; i < bones.size(); ++i)
	{
		auto & b = bones[i];

		b.position += b.velocity * dt;
		b.direction = geom::normalized(b.direction + geom::ort(b.direction) * b.angular_velocity * dt);
	}

	std::vector<constraint<1, 1>> constraints_1_1;
	std::vector<constraint<1, 2>> constraints_1_2;
	std::vector<constraint<2, 1>> constraints_2_1;
	std::vector<constraint<2, 2>> constraints_2_2;

	if (sel)
	{
		auto & b = bones[sel->index];

		constraint<1, 2> & c = constraints_1_2.emplace_back();
		c.bone[0] = sel->index;

		c.jacobian[0][0] = 1.f;
		c.jacobian[0][1] = 0.f;
		c.jacobian[0][2] = sel->pos * b.length / 2.f;
		c.jacobian[0][3] = 0.f;

		c.jacobian[1][0] = 0.f;
		c.jacobian[1][1] = 1.f;
		c.jacobian[1][2] = 0.f;
		c.jacobian[1][3] = sel->pos * b.length / 2.f;

		c.bias = sel->delta * 0.5f / dt;
	}

	for (std::size_t i = 0; i < bones.size(); ++i)
	{
		auto & b = bones[i];

		float const depth_factor = 0.5f / dt;

		float const bounce = 0.05f;
		float const friction = 0.05f;

		auto push_1 = [&](float depth, float s, geom::vector<float, 2> const & v){
			constraint<1, 1> & c = constraints_1_1.emplace_back();

			c.bone[0] = i;

			c.jacobian[0][0] = 0.f;
			c.jacobian[0][1] = 1.f;
			c.jacobian[0][2] = 0.f;
			c.jacobian[0][3] = b.length / 2.f * s;

			c.bias[0] = v[1] - depth * depth_factor;

			c.range.min = 0.f;

			constraint<1, 1> & fc = constraints_1_1.emplace_back();

			fc.bone[0] = i;

			fc.jacobian[0][0] = 1.f;
			fc.jacobian[0][1] = 0.f;
			fc.jacobian[0][2] = s * b.length / 2.f;
			fc.jacobian[0][3] = 0.f;

			fc.bias = fc.bias.zero();

			fc.range.min = -friction;
			fc.range.max = friction;
		};

		std::optional<float> d0, d1;

		auto p0 = b.position - b.direction * b.length / 2.f;
		auto p1 = b.position + b.direction * b.length / 2.f;

		auto v0 = b.velocity - geom::ort(b.direction) * b.angular_velocity * b.length / 2.f;
		auto v1 = b.velocity + geom::ort(b.direction) * b.angular_velocity * b.length / 2.f;

		if (p0[1] < b.width / 2.f)
		{
			float d = b.width / 2.f - p0[1];
			if (d > 0.f)
			{
				d0 = d;
			}
		}

		if (p1[1] < b.width / 2.f)
		{
			float d = b.width / 2.f - p1[1];
			if (d > 0.f)
			{
				d1 = d;
			}
		}

		if (d0 && d1)
		{
			constraint<1, 2> & c = constraints_1_2.emplace_back();

			c.bone[0] = i;

			c.jacobian[0][0] = 0.f;
			c.jacobian[0][1] = 1.f;
			c.jacobian[0][2] = 0.f;
			c.jacobian[0][3] = -b.length / 2.f;
			c.jacobian[1][0] = 0.f;
			c.jacobian[1][1] = 1.f;
			c.jacobian[1][2] = 0.f;
			c.jacobian[1][3] = b.length / 2.f;

			c.bias[0] = - (*d0) * depth_factor + v0[1] * bounce;
			c.bias[1] = - (*d1) * depth_factor + v1[1] * bounce;

			c.range.min = 0.f;

			constraint<1, 2> & fc = constraints_1_2.emplace_back();

			fc.bone[0] = i;

			fc.jacobian[0][0] = 1.f;
			fc.jacobian[0][1] = 0.f;
			fc.jacobian[0][2] = -b.length / 2.f;
			fc.jacobian[0][3] = 0.f;
			fc.jacobian[1][0] = 1.f;
			fc.jacobian[1][1] = 0.f;
			fc.jacobian[1][2] = b.length / 2.f;
			fc.jacobian[1][3] = 0.f;

			fc.range.min = -friction;
			fc.range.max = friction;

			fc.bias = fc.bias.zero();

		}
		else if (d0)
		{
			push_1(*d0, -1.f, v0 * bounce);
		}
		else if (d1)
		{
			push_1(*d1, 1.f, v1 * bounce);
		}
	}

	for (auto const & j : joints)
	{
		constraint<2, 2> & c = constraints_2_2.emplace_back();

		auto & b0 = bones[j.i0];
		auto & b1 = bones[j.i1];

		c.bone[0] = j.i0;
		c.bone[1] = j.i1;

		c.jacobian = c.jacobian.zero();
		c.jacobian[0][0] = -1.f;
		c.jacobian[0][2] = -j.s0 * b0.length / 2.f;
		c.jacobian[0][4] = 1.f;
		c.jacobian[0][6] = j.s1 * b1.length / 2.f;
		c.jacobian[1][1] = -1.f;
		c.jacobian[1][3] = -j.s0 * b0.length / 2.f;
		c.jacobian[1][5] = 1.f;
		c.jacobian[1][7] = j.s1 * b1.length / 2.f;

		c.bias = ((b1.position + b1.direction * j.s1 * b1.length / 2.f) - (b0.position + b0.direction * j.s0 * b0.length / 2.f)) / dt;

		if (j.angle_range)
		{
			float x = geom::dot(b0.direction, b1.direction);
			float y = geom::det(b0.direction, b1.direction);

			// a \in [-pi, pi]
			float a = std::atan2(y, x);

			// j.angle_range->max \in [0, 2*pi]
			// j.angle_range->min \in [-2*pi, 2*pi]
			// j.angle_range->center() \in [-pi, 2*pi]
			// delta \in [-3*pi, 2*pi]
			float delta = a - j.angle_range->center();
			if (delta > geom::pi)
				delta -= 2.f * geom::pi;
			if (delta < -geom::pi)
				delta += 2.f * geom::pi;
			// delta in [-pi, pi]

			float const half_length = j.angle_range->length() / 2.f;

			if (std::abs(delta) > half_length)
			{
				float const factor = 2.f;

				constraint<2, 1> & c = constraints_2_1.emplace_back();

				c.bone[0] = j.i0;
				c.bone[1] = j.i1;

				c.jacobian = c.jacobian.zero();
				c.jacobian[0][2] = - y * b1.direction[0] + x * b1.direction[1];
				c.jacobian[0][3] = - y * b1.direction[1] - x * b1.direction[0];
				c.jacobian[0][6] = - y * b0.direction[0] - x * b0.direction[1];
				c.jacobian[0][7] = - y * b0.direction[1] + x * b0.direction[0];

				if (delta > 0.f)
				{
					c.bias[0] = (delta - half_length) * factor / dt;
					c.range.max = 0.f;
				}
				else
				{
					c.bias[0] = (delta + half_length) * factor / dt;
					c.range.min = 0.f;
				}
			}
		}
	}

	std::size_t const steps = 1;

	for (std::size_t step = 0; step < steps; ++step)
	{
		for (auto const & c : constraints_1_1)
		{
			solve_constraint(c);
		}

		for (auto const & c : constraints_1_2)
		{
			solve_constraint(c);
		}

		for (auto const & c : constraints_2_1)
		{
			solve_constraint(c);
		}

		for (auto const & c : constraints_2_2)
		{
			solve_constraint(c);
		}
	}
}

template <std::size_t B, std::size_t C>
void system::solve_constraint(constraint<B, C> const & c)
{
	geom::vector<float, 4 * B> v;
	auto j_inv_mass = c.jacobian;

	for (std::size_t i = 0; i < B; ++i)
	{
		bone const & b = bones[c.bone[i]];

		v[4 * i + 0] = b.velocity[0];
		v[4 * i + 1] = b.velocity[1];
		v[4 * i + 2] = -b.angular_velocity * b.direction[1];
		v[4 * i + 3] = b.angular_velocity * b.direction[0];

		for (std::size_t j = 0; j < C; ++j)
		{
			j_inv_mass[j][4 * i + 0] /= b.mass;
			j_inv_mass[j][4 * i + 1] /= b.mass;
			j_inv_mass[j][4 * i + 2] /= b.inertia;
			j_inv_mass[j][4 * i + 3] /= b.inertia;
		}
	}

	geom::vector<float, C> l = *geom::solve(j_inv_mass * geom::transpose(c.jacobian), - c.jacobian * v - c.bias);
	for (std::size_t j = 0; j < C; ++j)
		l[j] = geom::clamp(l[j], c.range);

	auto p = geom::transpose(j_inv_mass) * l;

	for (std::size_t i = 0; i < B; ++i)
	{
		bone  & b = bones[c.bone[i]];

		b.velocity[0] += p[4 * i + 0];
		b.velocity[1] += p[4 * i + 1];
		b.angular_velocity += geom::det(b.direction, geom::vector{p[4 * i + 2], p[4 * i + 3]});
	}
}

struct controller
{
	// inputs should be equal to number of bones * 7
	// outputs should be equal to number of joints

	static constexpr std::size_t inputs = 24 * 7;
	static constexpr std::size_t outputs = 23;

//	static constexpr std::size_t layer1 = outputs;
//	static constexpr std::size_t layer2 = 4;
//	static constexpr std::size_t layer3 = outputs;

	//static constexpr std::size_t param_count = (inputs + 1) * outputs;

	static constexpr float max_output = 20.f;

	geom::matrix<float, outputs, inputs> m1;
	geom::vector<float, outputs> t1;

//	geom::matrix<float, layer1, inputs> m1;
//	geom::vector<float, layer1> t1;
//	geom::matrix<float, layer2, layer1> m2;
//	geom::vector<float, layer2> t2;
//	geom::matrix<float, layer3, layer2> m3;
//	geom::vector<float, layer3> t3;

	//Eigen::VectorXf to_eigen() const;
	//void from_eigen(Eigen::VectorXf const & v);

	int generation = 0;

	template <typename RNG>
	void randomize(RNG && rng, float amplitude);

	template <typename RNG>
	void mutate(RNG && rng, float amplitude);

	void reset();

	geom::vector<float, outputs> apply(system const & s) const;

	void read(std::istream & is);
	void write(std::ostream & os) const;

	static float activation(float x);

	template <std::size_t N>
	static geom::vector<float, N> activation(geom::vector<float, N> v);

	template <typename T>
	static void read(std::istream & is, T & x);

	template <typename T>
	static void write(std::ostream & os, T const & x);
};

template <typename RNG>
void controller::randomize(RNG && rng, float amplitude)
{
	generation = 0;

	random::uniform_distribution<float> d{-1.f, 1.f};

	auto visit_m = [&rng, &d, amplitude](auto & m)
	{
		for (std::size_t i = 0; i < m.rows(); ++i)
		{
			for (std::size_t j = 0; j < m.columns(); ++j)
			{
				m[i][j] = d(rng) * amplitude;
			}
		}
	};

	auto visit_v = [&rng, &d, amplitude](auto & v)
	{
		for (std::size_t i = 0; i < v.dimension(); ++i)
		{
			v[i] = d(rng) * amplitude;
		}
	};

	visit_m(m1);
//	visit_m(m2);
//	visit_m(m3);
	visit_v(t1);
//	visit_v(t2);
//	visit_v(t3);
}

/*
Eigen::VectorXf controller::to_eigen() const
{
	Eigen::VectorXf vec;
	vec.setZero(param_count);
	auto p = std::copy(m1.coords, m1.coords + inputs * outputs, vec.data());
	std::copy(t1.coords, t1.coords + outputs, p);
	return vec;
}

void controller::from_eigen(Eigen::VectorXf const & v)
{
	if (v.size() != param_count)
		throw std::runtime_error("Wrong data size");

	std::copy(v.data(), v.data() + inputs * outputs, m1.coords);
	std::copy(v.data() + inputs * outputs, v.data() + (inputs + 1) * outputs, t1.coords);
}
*/

template <typename RNG>
void controller::mutate(RNG && rng, float amplitude)
{
	random::normal_distribution<float> d{};

	auto visit_m = [&rng, &d, amplitude](auto & m)
	{
		for (std::size_t i = 0; i < m.rows(); ++i)
		{
			for (std::size_t j = 0; j < m.columns(); ++j)
			{
				m[i][j] += d(rng) * amplitude;
			}
		}
	};

	auto visit_v = [&rng, &d, amplitude](auto & v)
	{
		for (std::size_t i = 0; i < v.dimension(); ++i)
		{
			v[i] += d(rng) * amplitude;
		}
	};

	visit_m(m1);
//	visit_m(m2);
//	visit_m(m3);
	visit_v(t1);
//	visit_v(t2);
//	visit_v(t3);
}

void controller::reset()
{}

geom::vector<float, controller::outputs> controller::apply(system const & sys) const
{
	if (sys.bones.size() * 7 != inputs)
		throw std::runtime_error(util::to_string("Wrong number of inputs: should be ", sys.bones.size() * 7));
	if (sys.joints.size() != outputs)
		throw std::runtime_error(util::to_string("Wrong number of outputs: should be ", sys.joints.size()));

	geom::vector<float, inputs> input;
	for (std::size_t i = 0; i < inputs / 7; ++i)
	{
		input[7 * i + 0] = sys.bones[i].position[0] - sys.bones[0].position[0];
		input[7 * i + 1] = sys.bones[i].position[1];
		input[7 * i + 2] = sys.bones[i].direction[0];
		input[7 * i + 3] = sys.bones[i].direction[1];
		input[7 * i + 4] = sys.bones[i].velocity[0];
		input[7 * i + 5] = sys.bones[i].velocity[1];
		input[7 * i + 6] = sys.bones[i].angular_velocity;
	}

	auto r1 = activation(m1 * input + t1);
//	auto r2 = activation(m2 * r1 + t2);
//	auto r3 = activation(m3 * r2 + t3);

	for (std::size_t i = 0; i < sys.joints.size(); ++i)
//		r3[i] *= 0.5f * (sys.bones[sys.joints[i].i0].mass + sys.bones[sys.joints[i].i1].mass);
		r1[i] *= 0.5f * (sys.bones[sys.joints[i].i0].mass + sys.bones[sys.joints[i].i1].mass);

	return r1 * max_output;
}

template <typename T>
void controller::read(std::istream & is, T & x)
{
	is.read(reinterpret_cast<char *>(&x), sizeof(x));
}

template <typename T>
void controller::write(std::ostream & os, T const & x)
{
	os.write(reinterpret_cast<char const *>(&x), sizeof(x));
}

void controller::read(std::istream & is)
{
	read(is, m1);
	read(is, t1);
//	read(is, m2);
//	read(is, t2);
//	read(is, m3);
//	read(is, t3);
}

void controller::write(std::ostream & os) const
{
	write(os, m1);
	write(os, t1);
//	write(os, m2);
//	write(os, t2);
//	write(os, m3);
//	write(os, t3);
}

float controller::activation(float x)
{
	return 2.f / (1.f + std::exp(- 2.f * x)) - 1.f;
}

template <std::size_t N>
geom::vector<float, N> controller::activation(geom::vector<float, N> v)
{
	for (std::size_t i = 0; i < N; ++i)
		v[i] = activation(v[i]);
	return v;
}

controller lerp(controller const & c1, controller const & c2, float t)
{
	controller c;
	c.m1 = geom::lerp(c1.m1, c2.m1, t);
//	c.m2 = geom::lerp(c1.m2, c2.m2, t);
//	c.m3 = geom::lerp(c1.m3, c2.m3, t);
	c.t1 = geom::lerp(c1.t1, c2.t1, t);
//	c.t2 = geom::lerp(c1.t2, c2.t2, t);
//	c.t3 = geom::lerp(c1.t3, c2.t3, t);
	c.generation = std::max(c1.generation, c2.generation) + 1;
	return c;
}

struct animation_2d_app
	: app::app
{
	animation_2d_app();

	void on_resize(int width, int height) override;
	void on_mouse_wheel(int delta) override;

	void on_key_down(SDL_Keycode key) override;

	void update() override;
	void present() override;

	void update_camera();

	geom::box<float, 2> view_bbox;
	bool centered = false;

	gfx::painter painter;

	util::clock<> frame_clock;
	util::clock<> test_clock;

	system physics;

	std::optional<std::size_t> selected;
	float selected_s = 0.f;

	random::generator rng{random::random_device{}};

	async::threadpool bg{"bg"};

	enum class mode
	{
		train,
		test,
	} mode = mode::train;

	std::vector<controller> population;
	std::size_t const population_size = 1024;
	std::size_t const max_train_frames = 10.f / physics.dt;
	std::size_t const max_train_variations = 1;
	float const position_variation_amplitude = 0.f;
	float const angle_variation_amplitude = 0.f;

	std::size_t train_iterations = 0;
	std::size_t const max_train_iterations = 1024*8;
	float const initial_variance = 10.f;
	static constexpr auto mutation_amplitude = [](float t){ return std::pow(10.f, 1.f + geom::lerp(0.f, -2.f, t)); };

	//Eigen::VectorXf mean;
	//Eigen::MatrixXf covariance;

	float best_score = 0.f;
	bool const warm_start = false;
	bool const enable_testing = true;
	bool testing_control = true;
	std::string const cache_location = "./runner";

	std::size_t test_id = 0;
	std::vector<float> test_speeds;

	void reset_state(system & sys) const;

	float eval_score(controller const & c, random::generator rng) const;
	void do_train();
	void do_optimize();
	void do_test();
};

animation_2d_app::animation_2d_app()
	: app("Animation 2D")
{
	view_bbox[1] = {-1.f, 15.f};

	vsync(false);
}

void animation_2d_app::on_resize(int width, int height)
{
	app::on_resize(width, height);

	update_camera();
}

void animation_2d_app::on_mouse_wheel(int delta)
{
	float p = std::pow(0.8f, delta);
	view_bbox[1].max *= p;

	update_camera();
}

void animation_2d_app::update_camera()
{
	float ratio = static_cast<float>(width()) / height();

	float cx = 0.f;
	if (centered)
	{
		float m = 0.f;
		for (auto const & b : physics.bones)
		{
			cx += b.position[0] * b.mass;
			m += b.mass;
		}
		cx /= m;
	}
	else
	{
		cx = ratio * view_bbox[1].length() / 2.f;
	}

	view_bbox[0] = geom::expand(geom::interval<float>::singleton(cx), ratio * view_bbox[1].length() / 2.f);
}

void animation_2d_app::on_key_down(SDL_Keycode key)
{
	if (key == SDLK_c)
	{
		centered = !centered;
		update_camera();
	}

	if (key == SDLK_p)
	{
		testing_control = !testing_control;
	}

	if (mode == mode::train)
	{
		if (key == SDLK_s)
		{
			train_iterations = max_train_iterations;
		}
	}
	if (mode == mode::test)
	{
		bool reset = false;
		if (key == SDLK_LEFT)
		{
			reset = true;
			test_id = (test_id + population.size() - 1) % population.size();
		}
		if (key == SDLK_RIGHT)
		{
			reset = true;
			test_id = (test_id + 1) % population.size();
		}

		if (reset)
		{
			reset_state(physics);
			test_clock.restart();

			float shiftx = random::uniform_distribution<float>{-1.f, 1.f}(rng) * position_variation_amplitude;
			float shifty = random::uniform_distribution<float>{0.f, 1.f}(rng) * position_variation_amplitude;
			float angle = random::uniform_distribution<float>{-1.f, 1.f}(rng) * geom::rad(angle_variation_amplitude);

			geom::plane_rotation<float, 2> rot{0, 1, angle};

			float miny = 0.f;
			for (auto & b : physics.bones)
			{
				b.position = rot(b.position) + geom::vector{shiftx, shifty};
				b.direction = rot(b.direction);
				miny = std::min(miny, b.position[1] - std::abs(b.direction[1]) * b.length / 2.f - b.width / 2.f);
			}
			for (auto & b : physics.bones)
				b.position[1] -= miny;

			population[test_id].reset();
			test_speeds.clear();
		}
	}
}

void animation_2d_app::update()
{
	if (mode == mode::train)
	{
		do_train();
		if (train_iterations >= max_train_iterations)
		{
			mode = mode::test;
			frame_clock.restart();
			test_clock.restart();
			reset_state(physics);

			std::ofstream os(cache_location, std::ios::binary);
			population[0].write(os);

			population[0].reset();

			vsync(true);
		}
	}
	else if (mode == mode::test)
	{
		if (!enable_testing)
		{
			stop();
			return;
		}

		do_test();
	}
}

void animation_2d_app::reset_state(system & sys) const
{
	sys.bones.clear();
	sys.joints.clear();

/*
	// Jumper
	sys.bones.push_back(bone{{0.5f, 0.f}, {1.f, 0.f}, {0.f, 0.f}, 0.f, 1.f, 0.25f});
	sys.bones.push_back(bone{{0.25f, 0.5f}, {0.5f, 1.f}, {0.f, 0.f}, 0.f, std::sqrt(1.25f), 0.25f});
	sys.bones.push_back(bone{{0.75f, 1.5f}, {0.5f, 1.f}, {0.f, 0.f}, 0.f, std::sqrt(1.25f), 0.25f});
	sys.joints.push_back(joint{0, 1, -1.f, -1.f, geom::interval<float>{geom::rad(60.f), geom::rad(120.f)}});
	sys.joints.push_back(joint{1, 2, 1.f, -1.f, geom::interval<float>{geom::rad(-30.f), geom::rad(30.f)}});
//*/

/*
	// Dachshund
	float s = 0.5f;
	float h = 0.3f;
	sys.bones.push_back(bone{{  -s, h*0.5f}, {0.f,  1.f}, {0.f, 0.f}, 0.f,   h, 0.25f});
	sys.bones.push_back(bone{{  -s, h*1.5f}, {0.f,  1.f}, {0.f, 0.f}, 0.f,   h, 0.25f});
	sys.bones.push_back(bone{{ 0.f, h*2.f }, {1.f,  0.f}, {0.f, 0.f}, 0.f, 2.f, 0.25f});
	sys.bones.push_back(bone{{   s, h*1.5f}, {0.f, -1.f}, {0.f, 0.f}, 0.f,   h, 0.25f});
	sys.bones.push_back(bone{{   s, h*0.5f}, {0.f, -1.f}, {0.f, 0.f}, 0.f,   h, 0.25f});
	sys.joints.push_back(joint{0, 1, 1.f, -1.f, geom::interval{geom::rad(  0.f), geom::rad(45.f)}});
	sys.joints.push_back(joint{1, 2, 1.f,   -s, geom::interval{geom::rad(240.f), geom::rad(300.f)}});
	sys.joints.push_back(joint{2, 3,   s, -1.f, geom::interval{geom::rad(240.f), geom::rad(300.f)}});
	sys.joints.push_back(joint{3, 4, 1.f, -1.f, geom::interval{geom::rad(  0.f), geom::rad(45.f)}});
//*/

/*
	// Cube
	sys.bones.push_back(bone{{0.5f, 0.f}, {1.f, 0.f}, {0.f, 0.f}, 0.f, 1.f, 0.25f});
	sys.bones.push_back(bone{{1.f, 0.5f}, {0.f, 1.f}, {0.f, 0.f}, 0.f, 1.f, 0.25f});
	sys.bones.push_back(bone{{0.5f, 1.f}, {-1.f, 0.f}, {0.f, 0.f}, 0.f, 1.f, 0.25f});
	sys.bones.push_back(bone{{0.f, 0.5f}, {0.f, -1.f}, {0.f, 0.f}, 0.f, 1.f, 0.25f});
	sys.joints.push_back(joint{0, 1, 1.f, -1.f, geom::interval<float>{geom::rad(60.f), geom::rad(120.f)}});
	sys.joints.push_back(joint{1, 2, 1.f, -1.f, geom::interval<float>{geom::rad(60.f), geom::rad(120.f)}});
	sys.joints.push_back(joint{2, 3, 1.f, -1.f, geom::interval<float>{geom::rad(60.f), geom::rad(120.f)}});
	sys.joints.push_back(joint{3, 0, 1.f, -1.f, geom::interval<float>{geom::rad(60.f), geom::rad(120.f)}});
//*/

/*
	// Cube with tiny legs
	sys.bones.push_back(bone{{0.5f, 0.f}, {1.f, 0.f}, {0.f, 0.f}, 0.f, 1.f, 0.25f});
	sys.bones.push_back(bone{{1.f, 0.5f}, {0.f, 1.f}, {0.f, 0.f}, 0.f, 1.f, 0.25f});
	sys.bones.push_back(bone{{0.5f, 1.f}, {-1.f, 0.f}, {0.f, 0.f}, 0.f, 1.f, 0.25f});
	sys.bones.push_back(bone{{0.f, 0.5f}, {0.f, -1.f}, {0.f, 0.f}, 0.f, 1.f, 0.25f});
	sys.joints.push_back(joint{0, 1, 1.f, -1.f, geom::interval<float>{geom::rad(60.f), geom::rad(120.f)}});
	sys.joints.push_back(joint{1, 2, 1.f, -1.f, geom::interval<float>{geom::rad(60.f), geom::rad(120.f)}});
	sys.joints.push_back(joint{2, 3, 1.f, -1.f, geom::interval<float>{geom::rad(60.f), geom::rad(120.f)}});
	sys.joints.push_back(joint{3, 0, 1.f, -1.f, geom::interval<float>{geom::rad(60.f), geom::rad(120.f)}});
//*/

/*
	// Roller
	int n = 2;
	for (int i = 0; i < n; ++i)
		sys.bones.push_back(bone{{i + 0.5f, 0.f}, {1.f, 0.f}, {0.f, 0.f}, 0.f, 1.f, 0.25f});
	sys.bones.push_back(bone{{n + 0.25f, std::sqrt(3.f) * 0.25f}, {0.5f, std::sqrt(3.f) * 0.5}, {0.f, 0.f}, 0.f, 1.f, 0.25f});
	sys.bones.push_back(bone{{n + 0.25f, std::sqrt(3.f) * 0.75f}, {-0.5f, std::sqrt(3.f) * 0.5}, {0.f, 0.f}, 0.f, 1.f, 0.25f});
	for (int i = 0; i < n; ++i)
		sys.bones.push_back(bone{{n - i - 0.5f, std::sqrt(3.f)}, {-1.f, 0.f}, {0.f, 0.f}, 0.f, 1.f, 0.25f});
	sys.bones.push_back(bone{{-0.25f, std::sqrt(3.f) * 0.75f}, {-0.5f, -std::sqrt(3.f) * 0.5}, {0.f, 0.f}, 0.f, 1.f, 0.25f});
	sys.bones.push_back(bone{{-0.25f, std::sqrt(3.f) * 0.25f}, {0.5f, -std::sqrt(3.f) * 0.5}, {0.f, 0.f}, 0.f, 1.f, 0.25f});

	for (int i = 0; i < 2 * n + 4; ++i)
		sys.joints.push_back(joint{i, (i + 1) % (2 * n + 4), 1.f, -1.f, geom::interval<float>{geom::rad(0.f), geom::rad(60.f)}});
//*/

/*
	// Worm
	int n = 5;
	for (int i = 0; i < n; ++i)
		sys.bones.push_back(bone{{i + 0.5f, 0.f}, {1.f, 0.f}, {0.f, 0.f}, 0.f, 1.f, 0.25f});
	for (int i = 0; i + 1 < n; ++i)
		sys.joints.push_back(joint{i, i + 1, 1.f, -1.f, geom::interval<float>{geom::rad(-30.f), geom::rad(30.f)}});
//*/


	// Caterpillar
	int n = 8;
	for (int i = 0; i < n; ++i)
		sys.bones.push_back(bone{{i + 0.5f, 0.5f}, {1.f, 0.f}, {0.f, 0.f}, 0.f, 1.f, 0.25f});
	for (int i = 0; i < n; ++i)
	{
		sys.bones.push_back(bone{{i + 0.25f, 0.25f}, {0.f, -1.f}, {0.f, 0.f}, 0.f, 0.5f, 0.25f});
		sys.bones.push_back(bone{{i + 0.25f, 0.75f}, {0.f, -1.f}, {0.f, 0.f}, 0.f, 0.5f, 0.25f});
	}

	for (int i = 0; i + 1 < n; ++i)
		sys.joints.push_back(joint{i, i + 1, 1.f, -1.f, geom::interval<float>{geom::rad(-30.f), geom::rad(30.f)}});
	for (int i = 0; i < n; ++i)
	{
		sys.joints.push_back(joint{i, n + 2 * i + 0, -0.5f, -1.f, geom::interval<float>{geom::rad(240.f), geom::rad(300.f)}});
		sys.joints.push_back(joint{i, n + 2 * i + 1,  0.5f, -1.f, geom::interval<float>{geom::rad(240.f), geom::rad(300.f)}});
	}
//*/

/*
	// Runner-1
	sys.bones.push_back(bone{{0.f, 0.5f}, {0.f, 1.f}, {0.f, 0.f}, 0.f, 1.f, 0.25f});
	sys.bones.push_back(bone{{0.5f, 1.f}, {1.f, 0.f}, {0.f, 0.f}, 0.f, 1.f, 0.25f});
	sys.bones.push_back(bone{{1.5f, 1.f}, {1.f, 0.f}, {0.f, 0.f}, 0.f, 1.f, 0.25f});
	sys.bones.push_back(bone{{2.f, 0.5f}, {0.f, -1.f}, {0.f, 0.f}, 0.f, 1.f, 0.25f});
	sys.joints.push_back(joint{0, 1, 1.f, -1.f, geom::interval{geom::rad(240.f), geom::rad(300.f)}});
	sys.joints.push_back(joint{1, 2, 1.f, -1.f, geom::interval{geom::rad(-15.f), geom::rad( 15.f)}});
	sys.joints.push_back(joint{2, 3, 1.f, -1.f, geom::interval{geom::rad(240.f), geom::rad(300.f)}});
//*/

/*
	// Runner: dog
	sys.bones.push_back(bone{{-1.f, 0.25f}, {0.f, 1.f}, {0.f, 0.f}, 0.f, 0.5f, 0.25f});
	sys.bones.push_back(bone{{-1.f, 0.75f}, {0.f, 1.f}, {0.f, 0.f}, 0.f, 0.5f, 0.25f});
	sys.bones.push_back(bone{{ 0.f, 1.f  }, {1.f, 0.f}, {0.f, 0.f}, 0.f, 2.f, 0.25f});
	sys.bones.push_back(bone{{ 1.f, 0.75f}, {0.f, -1.f}, {0.f, 0.f}, 0.f, 0.5f, 0.25f});
	sys.bones.push_back(bone{{ 1.f, 0.25f}, {0.f, -1.f}, {0.f, 0.f}, 0.f, 0.5f, 0.25f});
	sys.joints.push_back(joint{0, 1, 1.f, -1.f, geom::interval{geom::rad(  0.f), geom::rad(45.f)}});
	sys.joints.push_back(joint{1, 2, 1.f, -1.f, geom::interval{geom::rad(240.f), geom::rad(300.f)}});
	sys.joints.push_back(joint{2, 3, 1.f, -1.f, geom::interval{geom::rad(240.f), geom::rad(300.f)}});
	sys.joints.push_back(joint{3, 4, 1.f, -1.f, geom::interval{geom::rad(  0.f), geom::rad(45.f)}});
//*/

/*
	// Runner: goat
	sys.bones.push_back(bone{{-1.f, 0.25f}, {0.f, 1.f}, {0.f, 0.f}, 0.f, 0.5f, 0.25f});
	sys.bones.push_back(bone{{-1.f, 0.75f}, {0.f, 1.f}, {0.f, 0.f}, 0.f, 0.5f, 0.25f});
	sys.bones.push_back(bone{{ 0.f, 1.f  }, {1.f, 0.f}, {0.f, 0.f}, 0.f, 2.f, 0.25f});
	sys.bones.push_back(bone{{ 1.f, 0.75f}, {0.f, -1.f}, {0.f, 0.f}, 0.f, 0.5f, 0.25f});
	sys.bones.push_back(bone{{ 1.f, 0.25f}, {0.f, -1.f}, {0.f, 0.f}, 0.f, 0.5f, 0.25f});
	sys.joints.push_back(joint{0, 1, 1.f, -1.f, geom::interval{geom::rad(-45.f), geom::rad(  0.f)}});
	sys.joints.push_back(joint{1, 2, 1.f, -1.f, geom::interval{geom::rad(240.f), geom::rad(300.f)}});
	sys.joints.push_back(joint{2, 3, 1.f, -1.f, geom::interval{geom::rad(240.f), geom::rad(300.f)}});
	sys.joints.push_back(joint{3, 4, 1.f, -1.f, geom::interval{geom::rad(-45.f), geom::rad(  0.f)}});
//*/

/*
	// Runner: left
	sys.bones.push_back(bone{{-1.f, 0.25f}, {0.f, 1.f}, {0.f, 0.f}, 0.f, 0.5f, 0.25f});
	sys.bones.push_back(bone{{-1.f, 0.75f}, {0.f, 1.f}, {0.f, 0.f}, 0.f, 0.5f, 0.25f});
	sys.bones.push_back(bone{{ 0.f, 1.f  }, {1.f, 0.f}, {0.f, 0.f}, 0.f, 2.f, 0.25f});
	sys.bones.push_back(bone{{ 1.f, 0.75f}, {0.f, -1.f}, {0.f, 0.f}, 0.f, 0.5f, 0.25f});
	sys.bones.push_back(bone{{ 1.f, 0.25f}, {0.f, -1.f}, {0.f, 0.f}, 0.f, 0.5f, 0.25f});
	sys.joints.push_back(joint{0, 1, 1.f, -1.f, geom::interval{geom::rad(-45.f), geom::rad(  0.f)}});
	sys.joints.push_back(joint{1, 2, 1.f, -1.f, geom::interval{geom::rad(240.f), geom::rad(300.f)}});
	sys.joints.push_back(joint{2, 3, 1.f, -1.f, geom::interval{geom::rad(240.f), geom::rad(300.f)}});
	sys.joints.push_back(joint{3, 4, 1.f, -1.f, geom::interval{geom::rad(  0.f), geom::rad( 45.f)}});
//*/

/*
	// Runner: right
	sys.bones.push_back(bone{{-1.f, 0.25f}, {0.f, 1.f}, {0.f, 0.f}, 0.f, 0.5f, 0.25f});
	sys.bones.push_back(bone{{-1.f, 0.75f}, {0.f, 1.f}, {0.f, 0.f}, 0.f, 0.5f, 0.25f});
	sys.bones.push_back(bone{{ 0.f, 1.f  }, {1.f, 0.f}, {0.f, 0.f}, 0.f, 2.f, 0.25f});
	sys.bones.push_back(bone{{ 1.f, 0.75f}, {0.f, -1.f}, {0.f, 0.f}, 0.f, 0.5f, 0.25f});
	sys.bones.push_back(bone{{ 1.f, 0.25f}, {0.f, -1.f}, {0.f, 0.f}, 0.f, 0.5f, 0.25f});
	sys.joints.push_back(joint{0, 1, 1.f, -1.f, geom::interval{geom::rad(  0.f), geom::rad( 45.f)}});
	sys.joints.push_back(joint{1, 2, 1.f, -1.f, geom::interval{geom::rad(240.f), geom::rad(300.f)}});
	sys.joints.push_back(joint{2, 3, 1.f, -1.f, geom::interval{geom::rad(240.f), geom::rad(300.f)}});
	sys.joints.push_back(joint{3, 4, 1.f, -1.f, geom::interval{geom::rad(-45.f), geom::rad(  0.f)}});
//*/

	float const bone_density = 1.f;
	for (auto & b : sys.bones)
	{
		b.position[1] += 0.125f;
		b.direction = geom::normalized(b.direction);
		b.mass = b.length * b.width * bone_density;
		b.inertia = b.mass * (b.length * b.length + b.width * b.width) / 12.f;
	}
}

float animation_2d_app::eval_score(controller const & c, random::generator rng) const
{
	system physics;
	std::optional<system::selection> sel;

	float score = 0.f;

	for (std::size_t variation = 0; variation < max_train_variations; ++variation)
	{
		reset_state(physics);

		float shiftx = random::uniform_distribution<float>{-1.f, 1.f}(rng) * position_variation_amplitude;
		float shifty = random::uniform_distribution<float>{0.f, 1.f}(rng) * position_variation_amplitude;
		float angle = random::uniform_distribution<float>{-1.f, 1.f}(rng) * geom::rad(angle_variation_amplitude);

		geom::plane_rotation<float, 2> rot{0, 1, angle};

		float miny = 0.f;
		for (auto & b : physics.bones)
		{
			b.position = rot(b.position) + geom::vector{shiftx, shifty};
			b.direction = rot(b.direction);
			miny = std::min(miny, b.position[1] - std::abs(b.direction[1]) * b.length / 2.f - b.width / 2.f);
		}
		for (auto & b : physics.bones)
			b.position[1] -= miny;

		float const dt = physics.dt;

		bool failure = false;
		float cur_score = 0.f;
		float energy = 0.f;
		float reward = 0.f;

		static std::vector<std::pair<std::size_t, float>> hit_floor_points =
		{
//			{2,  1.f},
//			{2, -1.f},
//			{2,  1.f},
//			{3, -1.f},
//			{4, -1.f},
//			{5, -1.f},
		};

		std::size_t train_frames = max_train_frames * 1.0f;

		float time = 0.f;

		std::vector<float> switch_times;
		for (int i = 0; i < 10; ++i)
			switch_times.push_back(random::uniform_distribution<float>{0.f, train_frames * dt}(rng));
		std::sort(switch_times.begin(), switch_times.end());

		for (std::size_t frame = 0; frame < train_frames; ++frame)
		{
			bool moving = false;

			for (auto t : switch_times)
			{
				if (time < t) break;
				moving = !moving;
			}

			physics.advance(dt, sel, [&c, &energy, dt](system const & physics){
				std::vector<float> torque(physics.joints.size(), 0.f);
				auto ctrl = c.apply(physics);
				for (std::size_t i = 0; i < ctrl.dimension(); ++i)
					torque[i] = ctrl[i];
				energy += geom::length(ctrl) * dt;
				return torque;
			});

			for (auto f : hit_floor_points)
			{
				if (f.first >= physics.bones.size()) break;

				auto const & b = physics.bones[f.first];

				auto const p = b.position + f.second * b.direction * b.length / 2.f;

				if (p[1] <= b.width / 2.f)
				{
					failure = true;
					break;
				}
			}

//			penalty += geom::sqr((1.f - physics.bones[2].direction[0]) * dt);
//			float const target_speed = 2.f;
//			penalty += geom::sqr((physics.bones[2].velocity[0] - target_speed) / target_speed) * dt;
//			penalty -= physics.bones[2].velocity[0] * dt / physics.bones;

			auto cm_vel = geom::vector<float, 2>::zero();
			float mass = 0.f;
			for (auto const & b : physics.bones)
			{
				cm_vel += b.velocity * b.mass;
				mass += b.mass;
			}
			cm_vel /= mass;

			reward += (cm_vel[0]) * dt;

//			reward -= energy * 0.1f;

			if (failure)
			{
				cur_score = 0.f;
				break;
			}

			if (frame + 1 == max_train_frames)
			{
				float mean_pos_x = 0.f;
				float mass = 0.f;
				for (auto const & b : physics.bones)
				{
					mean_pos_x += b.position[0] * b.mass;
					mass += b.mass;
				}
				mean_pos_x /= mass;
				cur_score = reward / max_train_variations;
			}

			time += dt;
		}

		(void)failure;

		score += cur_score;
	}

	return score;
}


// Old evolutionary training implementation
void animation_2d_app::do_train()
{
	if (population.empty())
	{
		population.resize(population_size);
		std::ifstream is(cache_location, std::ios::binary);
		if (warm_start && is)
		{
			controller w;
			w.read(is);
			for (auto & c : population)
			{
				c = w;
			}
		}
		else
		{
			for (auto & c : population)
				c.randomize(rng, initial_variance);
		}
	}
	else
	{
		std::size_t const preserved = population.size() / 16;
		std::size_t const cross_parents = population.size() / 8;
		std::size_t const cross_new = population.size() / 2;

		std::vector<controller> new_population(population.size());

		for (std::size_t i = 0; i < preserved; ++i)
		{
			new_population[i] = population[i];
		}

		float const mu = mutation_amplitude((train_iterations * 1.f) / max_train_iterations);

		for (std::size_t i = preserved; i < preserved + cross_new; ++i)
		{
			auto p1 = random::uniform_distribution<std::size_t>{0, cross_parents - 1}(rng);
			auto p2 = random::uniform_distribution<std::size_t>{0, cross_parents - 1}(rng);

			float t = random::uniform_distribution<float>{}(rng);
			new_population[i] = lerp(population[p1], population[p2], t);
			new_population[i].mutate(rng, mu);
		}

		for (std::size_t i = preserved + cross_new; i < population.size(); ++i)
			new_population[i].randomize(rng, initial_variance);

		population = std::move(new_population);
	}

	//(void)&controller::to_eigen;
	//(void)&controller::from_eigen;

	std::vector<std::pair<float, std::size_t>> scores(population.size());
	std::vector<async::future<void>> futures;
	for (std::size_t i = 0; i < population.size(); ++i)
	{
		futures.push_back(bg.dispatch([&, i, rng = rng]() mutable {
			scores[i] = {eval_score(population[i], rng), i};
		}));
	}
	bg.wait_all(futures.begin(), futures.end()).get();

	std::sort(scores.begin(), scores.end(), [](auto const & p1, auto const & p2){ return p1.first > p2.first; });

	std::vector<controller> new_population(population.size());
	for (std::size_t i = 0; i < population.size(); ++i)
	{
		new_population[i] = population[scores[i].second];
	}
	population = std::move(new_population);
	++train_iterations;

	best_score = scores.front().first;
}


/*
// Covariance matrix adaptation
void animation_2d_app::do_train()
{
	if (population.empty())
	{
		population.resize(population_size);
		std::ifstream is(cache_location, std::ios::binary);
		if (warm_start && is)
		{
			controller w;
			w.read(is);
			for (auto & c : population)
			{
				c = w;
				c.mutate(rng, initial_variance);
			}
		}
		else
		{
			for (auto & c : population)
				c.randomize(rng, initial_variance);
		}

		mean.setZero(controller::param_count);
		for (auto const & c : population)
		{
			Eigen::VectorXf v = c.to_eigen();
//			mean += v;
		}
		mean /= static_cast<float>(population.size());

		covariance.setZero(controller::param_count, controller::param_count);
		for (auto const & c : population)
		{
			Eigen::VectorXf const v = c.to_eigen() - mean;
			covariance += v * v.transpose();
		}
		covariance /= static_cast<float>(population.size() - 1);
	}
	else
	{
#pragma GCC diagnostic push
#pragma GCC diagnostic ignored "-Wmaybe-uninitialized"
		auto ldlt = covariance.ldlt();
		Eigen::MatrixXf const l = ldlt.matrixL();
		Eigen::VectorXf const d = ldlt.vectorD();

		random::normal_distribution<float> n;

		Eigen::VectorXf v(controller::param_count);

		for (std::size_t i = 0; i < population.size(); ++i)
		{
			for (std::size_t j = 0; j < controller::param_count; ++j)
				v[j] = n(rng);

			for (std::size_t j = 0; j < controller::param_count; ++j)
			{
//				if (d[j] < 0.f) throw std::runtime_error(util::to_string("Covariance matrix is not positive-semidefinite: ", d[j]));
				v[j] *= std::sqrt(std::max(d[j], 0.f));
			}

			population[i].from_eigen(l * v + mean);
		}
#pragma GCC diagnostic pop
	}

	std::vector<std::pair<float, std::size_t>> scores(population.size());
	std::vector<async::future<void>> futures(population.size());
	for (std::size_t i = 0; i < population.size(); ++i)
	{
		futures[i] = bg.dispatch([&, i, rng = rng]() mutable {
			scores[i] = {eval_score(population[i], rng), i};
		});
	}
	bg.wait_all(futures.begin(), futures.end()).get();

	std::sort(scores.begin(), scores.end(), [](auto const & p1, auto const & p2){ return p1.first > p2.first; });

	std::size_t const selected = population.size() / 4;

	auto weight = [selected](std::size_t rank, float){
		return static_cast<float>(selected - rank);
	};

	float sum_weights = 0.f;
	for (std::size_t i = 0; i < selected; ++i)
		sum_weights += weight(scores[i].second, scores[i].first);

	mean.setZero(controller::param_count);
	for (std::size_t i = 0; i < selected; ++i)
	{
		float w = weight(scores[i].second, scores[i].first) / sum_weights;
		mean += w * population[scores[i].second].to_eigen();
	}

	Eigen::MatrixXf new_covariance;
	new_covariance.setZero(controller::param_count, controller::param_count);
	for (std::size_t i = 0; i < selected; ++i)
	{
		float w = weight(scores[i].second, scores[i].first) / sum_weights;
		Eigen::VectorXf const v = population[scores[i].second].to_eigen() - mean;
		new_covariance += w * v * v.transpose();
	}
//	new_covariance *= static_cast<float>(population.size() - 1) / static_cast<float>(population.size());

	float const mu = 0.0001f;

	covariance = covariance * (1.f - mu) + new_covariance * mu;

	best_score = scores.front().first;

	++train_iterations;

	(void)lerp;
}
*/

void animation_2d_app::do_test()
{
	std::optional<geom::point<float, 2>> m;
	if (mouse())
	{
		m = view_bbox.corner((*mouse())[0] * 1.f / width(), 1.f - (*mouse())[1] * 1.f / height());
	}

	if (!is_left_button_down())
	{
		selected = std::nullopt;

		if (m)
		{
			float selected_distance = std::numeric_limits<float>::infinity();

			for (std::size_t i = 0; i < physics.bones.size(); ++i)
			{
				auto & b = physics.bones[i];

				auto p0 = b.position - b.direction * b.length / 2.f;
				auto p1 = b.position + b.direction * b.length / 2.f;
				auto r = p1 - p0;

				auto d = *m - p0;

				float t = geom::dot(d, r) / geom::dot(r, r);

				float distance;

				if (0.f <= t && t <= 1.f)
				{
					distance = geom::length(d - r * t);
				}
				else
				{
					float d0 = geom::distance(p0, *m);
					float d1 = geom::distance(p1, *m);

					if (d0 < d1)
					{
						t = 0.f;
					}
					else
					{
						t = 1.f;
					}

					distance = std::min(d0, d1);
				}

				if (distance < b.width / 2.f && distance < selected_distance)
				{
					selected_distance = distance;
					selected = i;
					selected_s = 2.f * t - 1.f;
				}
			}
		}
	}

	std::optional<system::selection> sel;
	if (selected && is_left_button_down())
	{
		auto const & b = physics.bones[*selected];
		auto delta = b.position + b.direction * selected_s * b.length / 2.f - *m;
		sel = system::selection{*selected, selected_s, delta};
	}

	if (testing_control)
	{
		physics.advance(frame_clock.restart().count(), sel, [this](system const & physics){
			std::vector<float> torque(physics.joints.size());
			auto ctrl = population[test_id].apply(physics);
			for (std::size_t i = 0; i < ctrl.dimension(); ++i)
				torque[i] = ctrl[i];
			return torque;
		});
	}
	else
	{
		physics.advance(frame_clock.restart().count(), sel, [](system const & physics){
			return std::vector<float>(physics.joints.size(), 0.f);
		});
	}

	{
		auto cm_vel = geom::vector<float, 2>::zero();
		float mass = 0.f;
		for (auto const & b : physics.bones)
		{
			cm_vel += b.velocity * b.mass;
			mass += b.mass;
		}
		cm_vel /= mass;

		test_speeds.push_back(physics.bones[2].velocity[0]);
	}
}

void animation_2d_app::present()
{
	update_camera();

	gl::ClearColor(0.8f, 0.8f, 0.8f, 1.f);
	gl::Clear(gl::COLOR_BUFFER_BIT);

	{
		float ground_width = 0.1f;
		painter.line({view_bbox[0].min, -ground_width/2.f}, {view_bbox[0].max, -ground_width/2.f}, ground_width, gfx::black);
	}

	for (auto const & b : physics.bones)
	{
		auto c = gfx::dark(gfx::blue).as_color_rgba();
		painter.line(b.position - b.direction * b.length / 2.f, b.position + b.direction * b.length / 2.f, b.width, c);
	}

	if (selected)
	{
		auto c = gfx::blue.as_color_rgba();

		auto const & b = physics.bones[*selected];
		painter.line(b.position - b.direction * b.length / 2.f, b.position + b.direction * b.length / 2.f, b.width, c);
	}

	for (auto const & b : physics.bones)
	{
		for (float s : {-1.f, 1.f})
		{
			auto c = gfx::dark(gfx::red).as_color_rgba();
			auto p = b.position + s * b.direction * b.length / 2.f;
			if (p[1] < b.width / 2.f)
				c = gfx::yellow.as_color_rgba();
			painter.circle(p, b.width * 0.35f, c);
		}
	}

	painter.render(geom::orthographic_camera{view_bbox}.transform());

	float avg_speed = 0.f;

	if (mode == mode::test)
	{
		for (float s : test_speeds)
			avg_speed += s / test_speeds.size();

		int margin = 40;
		float const step = 1.f;

		int max_frames_shown = (width() - 2 * margin) / step;
		int start = std::max(0, static_cast<int>(test_speeds.size() - max_frames_shown));

		for (std::size_t i = start; i + 1 < test_speeds.size(); ++i)
		{
			float const scale = 2.f;
			geom::point p0{40.f + (i - start    ) * step, 180.f - test_speeds[i    ] * scale};
			geom::point p1{40.f + (i - start + 1) * step, 180.f - test_speeds[i + 1] * scale};
			painter.line(p0, p1, 2.f, gfx::red, false);
		}
	}

	{
		gfx::painter::text_options opts;
		opts.c = gfx::black;
		opts.x = gfx::painter::x_align::left;
		opts.y = gfx::painter::y_align::top;
		opts.scale = 2.f;
		painter.text({40.f, 40.f}, util::to_string(train_iterations, "/", max_train_iterations), opts);
		painter.text({40.f, 64.f}, util::to_string("Best score: ", std::setprecision(10), best_score), opts);
		painter.text({40.f, 88.f}, util::to_string("Model: ", test_id, "/", population.size(), ", gen ", population[test_id].generation), opts);
		painter.text({40.f, 112.f}, util::to_string(util::pretty(test_clock.duration(), std::chrono::milliseconds{1})), opts);
		painter.text({40.f, 136.f}, util::to_string("View height: ", view_bbox[1].length()), opts);

//		if (mode == mode::test && !test_speeds.empty()) painter.text({40.f, 136.f}, util::to_string("Speed: ", test_speeds.back()), opts);
	}

	painter.render(geom::window_camera{width(), height()}.transform());
}

}

int main()
{
	return app::main<animation_2d_app>();
}