#include <psemek/app/application_base.hpp>
#include <psemek/app/default_application_factory.hpp>
#include <psemek/gfx/painter.hpp>
#include <psemek/gfx/gl.hpp>
#include <psemek/math/camera.hpp>
#include <psemek/math/contains.hpp>
#include <psemek/cg/kdtree.hpp>
#include <psemek/random/device.hpp>
#include <psemek/random/generator.hpp>
#include <psemek/random/uniform.hpp>
#include <psemek/random/uniform_box.hpp>
#include <psemek/util/clock.hpp>
#include <psemek/util/array.hpp>
#include <psemek/log/log.hpp>

#include <format>

using namespace psemek;

struct particle
{
	math::point<float, 2> position;
	math::vector<float, 2> velocity;
	float mass;
	float radius;
	int type;
};

struct particle_life_2d_app
	: app::application_base
{
	particle_life_2d_app(options const &, context const &)
	{
		float S = 5.f;
		float W = 160.f * S;
		float H =  90.f * S;
		area_ = {{{-W, W}, {-H, H}}};

		types_ = 10;

		auto seed = random::device{}();

		// Explosions
		// types_ = 6;
		// seed = 3940378904;

		// Cells
		// types_ = 8;
		// seed = 3071915692;

		// Cells v2
		// types_ = 6;
		// seed = 3337663839;

		// Cool cells!
		types_ = 10;
		seed = 388834085;

		// One cell, use central gravity
		// types_ = 10;
		// seed = 175098945;

		// Mitochondria cell
		// types_ = 10;
		// seed = 1763331406;

		// Just fun
		// types_ = 10;
		// seed = 1676103531;

		// Another cell
		// types_ = 10;
		// seed = 2396049785;

		log::info() << "Seed: " << seed;

		random::generator rng{seed, 0};

		random::uniform_distribution<int> dcolor(63, 255);
		random::uniform_distribution<int> dtype(0, types_ - 1);
		random::uniform_box_point_distribution<float, 2> darea(area_);

		for (int i = 0; i < types_; ++i)
			colors_.push_back({dcolor(rng), dcolor(rng), dcolor(rng), 255});

		for (int i = 0; i < 4 * 1024; ++i)
		{
			particles_.push_back({
				.position = darea(rng),
				.velocity = {0.f, 0.f},
				.mass = 1.f,
				.radius = 1.f,
				.type = dtype(rng),
				// .type = random::uniform_from(rng, {4, 6}),
			});
		}

		auto dforce = [](auto & rng)
		{
			return random::uniform(rng, 20.f, 100.f) * (random::uniform<bool>(rng) ? 1.f : -1.f);
		};

		random::uniform_distribution<float> ddistance(3.f, 20.f);

		force_constants_.resize({types_, types_});
		force_distance_.resize({types_, types_});
		collision_distance_.resize({types_, types_});

		for (int i = 0; i < types_; ++i)
			for (int j = 0; j < types_; ++j)
				force_constants_(i, j) = dforce(rng);

		for (int i = 0; i < types_; ++i)
			for (int j = 0; j < types_; ++j)
				force_distance_(i, j) = ddistance(rng);

		for (int i = 0; i < types_; ++i)
			for (int j = 0; j < types_; ++j)
				collision_distance_(i, j) = random::uniform(rng, 0.125f, 0.75f) * force_distance_(i, j);

		// for (int i = 0; i < types_; ++i)
		// 	for (int j = 0; j < types_; ++j)
		// 		force_constants_(i, j) = -100.f * ((i == j) ? 1.f : 2.f);

		max_force_distance_.resize(types_, 2.f);
		for (int i = 0; i < types_; ++i)
			for (int j = 0; j < types_; ++j)
				math::make_max(max_force_distance_[i], force_distance_(i, j));
	}

	void on_event(app::mouse_wheel_event const & event) override
	{
		app::application_base::on_event(event);

		scale_target_ *= std::pow(1.25f, -event.delta);
	}

	void on_event(app::key_event const & event) override
	{
		app::application_base::on_event(event);

		if (event.down && event.key == app::keycode::SPACE)
			paused_ ^= true;
	}

	void update() override
	{
		float const dt = 0.02f;

		scale_ += (scale_target_ - scale_) * (- std::expm1(- 20.f * dt));

		auto visible_area = area_;
		visible_area = math::expand(visible_area, 0.5f * (scale_ - 1.f) * visible_area.dimensions());

		float window_aspect_ratio = state().size[0] * 1.f / state().size[1];
		float area_aspect_ratio = visible_area[0].length() / visible_area[1].length();

		if (area_aspect_ratio < window_aspect_ratio)
		{
			view_area_[1] = visible_area[1];
			view_area_[0] = math::expand(math::interval<float>::singleton(visible_area[0].center()), visible_area[1].length() * window_aspect_ratio * 0.5f);
		}
		else
		{
			view_area_[0] = visible_area[0];
			view_area_[1] = math::expand(math::interval<float>::singleton(visible_area[1].center()), visible_area[0].length() / window_aspect_ratio * 0.5f);
		}

		if (!paused_)
		for (int step = 0; step < 5; ++step)
		{
			cg::kdtree<float, 2, int> kdtree;

			for (int i = 0; i < particles_.size(); ++i)
				kdtree.insert({particles_[i].position, i});

			float const collision_strength = 200.f;

			mouseover_particle_ = std::nullopt;
			if (state().mouse_button_down.contains(app::mouse_button::left))
			{
				math::point<float, 2> m;
				m[0] = math::lerp(view_area_[0], state().mouse[0] * 1.f / state().size[0]);
				m[1] = math::lerp(view_area_[1], 1.f - state().mouse[1] * 1.f / state().size[1]);
				mouseover_particle_ = kdtree.closest(m).data;
			}

			#pragma omp parallel for
			for (int i = 0; i < particles_.size(); ++i)
			{
				// ======== kd-tree based algorithm ========

				auto & pi = particles_[i];
				auto f = math::vector{0.f, 0.f};

				kdtree.closer_than_map(pi.position, max_force_distance_[pi.type], [&](auto const & value){
					auto j = value.data;

					if (i == j)
						return;

					auto const & pj = particles_[j];

					auto r = pj.position - pi.position;
					float l = math::length(r);
					auto n = r / l;

					// ==== 1/R forces ====
					// float lg = std::max(0.5f, l);
					// auto g = n / (lg * lg);

					// if (l < force_distance_(pi.type, pj.type))
					// 	f += force_constants_(pi.type, pj.type) * g;

					// float collision_distance = pi.radius + pj.radius;

					// if (l < collision_distance)
					// {
					// 	float d = std::max(0.25f, l / collision_distance);

					// 	f -= collision_strength * n * (1.f / (d * d) - 1.f);
					// }

					// ==== Linear forces ====
					if (l < force_distance_(pi.type, pj.type))
						f += force_constants_(pi.type, pj.type) * n * (1.f - l / force_distance_(pi.type, pj.type));

					if (l < collision_distance_(pi.type, pj.type))
						f -= 10.f * std::abs(force_constants_(pi.type, pj.type)) * n * (1.f - l / collision_distance_(pi.type, pj.type));

					(void)collision_strength;
				});

				pi.velocity += f * dt / pi.mass;






				// ======== Naive algorithm v2 ========

				// auto & pi = particles_[i];
				// auto f = math::vector{0.f, 0.f};

				// for (int j = 0; j < particles_.size(); ++j)
				// {
				// 	if (i == j)
				// 		continue;

				// 	auto const & pj = particles_[j];

				// 	auto r = pj.position - pi.position;
				// 	float l = math::length(r);
				// 	auto n = r / l;
				// 	auto g = n / (l * l);

				// 	if (l < force_distance_(pi.type, pj.type))
				// 		f += force_constants_(pi.type, pj.type) * g;

				// 	if (l < pi.radius + pj.radius)
				// 	{
				// 		float d = l / (pi.radius + pj.radius);

				// 		f -= collision_strength * n * (1.f / (d * d) - 1.f);
				// 	}
				// }

				// pi.velocity += f * dt / pi.mass;





				// ======== Naive algorithm ========

				// for (int j = i + 1; j < particles_.size(); ++j)
				// {
				// 	auto & pi = particles_[i];
				// 	auto & pj = particles_[j];

				// 	auto fi = math::vector{0.f, 0.f};
				// 	auto fj = math::vector{0.f, 0.f};

				// 	auto r = pj.position - pi.position;
				// 	float l = math::length(r);
				// 	auto n = r / l;
				// 	auto g = n / (l * l);

				// 	if (l < force_distance_(pi.type, pj.type))
				// 		fi += force_constants_(pi.type, pj.type) * g;

				// 	if (l < force_distance_(pj.type, pi.type))
				// 		fj -= force_constants_(pj.type, pi.type) * g;

				// 	if (l < pi.radius + pj.radius)
				// 	{
				// 		float d = l / (pi.radius + pj.radius);

				// 		auto f = collision_strength * n * (1.f / (d * d) - 1.f);
				// 		fi -= f;
				// 		fj += f;
				// 	}

				// 	pi.velocity += fi * dt / pi.mass;
				// 	pj.velocity += fj * dt / pj.mass;
				// }
			}

			float const friction = 10.f;
			float const friction_factor = std::exp(- friction * dt);
			float const wall_distance = 0.f;
			float const wall_force = 10.f;
			float const gravity = 0.f;
			float const center_gravity = 0.f;
			float const line_gravity = 0.f;
			float const circle_gravity = 0.f;
			float const circle_radius = 100.f;
			bool const circle_wall = false;
			float const circle_wall_radius = 100.f;

			for (auto & p : particles_)
			{
				auto r = p.position - p.position.zero();

				p.velocity[1] -= gravity * dt;
				p.velocity[1] -= line_gravity * p.position[1] * dt;
				p.velocity -= center_gravity * r * dt;
				p.velocity += circle_gravity * r / math::length(r) * (circle_radius - math::length(r)) * dt;
				p.velocity *= friction_factor;
				p.position += p.velocity * dt;

				for (int d : {0, 1})
				{
					// if (p.position[d] < area_[d].min + p.radius)
					// {
					// 	p.position[d] = area_[d].min + p.radius;
					// 	p.velocity[d] *= -1.f;
					// }

					// if (p.position[d] > area_[d].max - p.radius)
					// {
					// 	p.position[d] = area_[d].max - p.radius;
					// 	p.velocity[d] *= -1.f;
					// }

					if (circle_wall)
					{
						auto r = p.position - p.position.zero();
						auto l = math::length(r);
						if (l > circle_wall_radius)
						{
							p.velocity -= wall_force * r / l * (l - circle_wall_radius) * dt;
						}
					}
					else
					{
						if (p.position[d] < area_[d].min + wall_distance)
						{
							p.velocity[d] += wall_force * (area_[d].min + wall_distance - p.position[d]) * dt;
						}

						if (p.position[d] > area_[d].max - wall_distance)
						{
							p.velocity[d] -= wall_force * (p.position[d] - (area_[d].max - wall_distance)) * dt;
						}
					}
				}
			}
		}
	}

	void present() override
	{
		gl::Viewport(0, 0, state().size[0], state().size[1]);

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

		for (auto const & p : particles_)
		{
			auto c0 = colors_[p.type];
			c0[3] = 15;

			auto c1 = colors_[p.type];
			c1[3] = 0;

			painter_.circle(p.position, p.radius * 8.f, c0, c1);
		}

		for (auto const & p : particles_)
			painter_.circle(p.position, p.radius, colors_[p.type]);

		painter_.render(math::orthographic_camera{view_area_}.transform());

		if (mouseover_particle_)
		{
			int type = particles_[*mouseover_particle_].type;
			painter_.text({state().size[0] / 2.f, state().size[1] - 20.f}, std::format("Species #{}", type), {
				.scale = {2.f, 2.f},
				.y = gfx::painter::y_align::bottom,
				.c = colors_[type],
			});
		}

		painter_.render(math::window_camera{state().size[0], state().size[1]}.transform());
	}

private:
	gfx::painter painter_;

	util::clock<std::chrono::duration<float>, std::chrono::high_resolution_clock> clock_;

	math::box<float, 2> area_;
	math::box<float, 2> view_area_;
	float scale_ = 1.f;
	float scale_target_ = 1.f;

	int types_;
	std::vector<gfx::color_rgba> colors_;
	util::array<float, 2> force_constants_;
	util::array<float, 2> force_distance_;
	util::array<float, 2> collision_distance_;
	std::vector<float> max_force_distance_;

	std::vector<particle> particles_;

	std::optional<int> mouseover_particle_;

	bool paused_ = true;
};

namespace psemek::app
{

	std::unique_ptr<application::factory> make_application_factory()
	{
		return default_application_factory<particle_life_2d_app>({.name = "Particle Life 2D"});
	}

}