#include <psemek/app/application_base.hpp>
#include <psemek/app/default_application_factory.hpp>
#include <psemek/gfx/painter.hpp>
#include <psemek/math/orthographic.hpp>
#include <psemek/math/camera.hpp>
#include <psemek/log/log.hpp>
#include <psemek/util/ndarray.hpp>
#include <psemek/random/device.hpp>
#include <psemek/random/generator.hpp>
#include <psemek/random/uniform.hpp>
#include <psemek/random/uniform_sphere.hpp>
#include <psemek/pcg/perlin.hpp>

#include <format>

using namespace psemek;

struct water_2d_app
	: app::application_base
{
	water_2d_app(options const &, context const &);

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

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

	math::box<float, 2> simulation_area;
	math::box<float, 2> inner_area;
	float aspect_ratio;

	random::generator rng{random::device{}};

	int const N = 256;

	float time = 0.f;

	util::ndarray<float, 2> bed;
	util::ndarray<float, 2> water;
	util::ndarray<float, 2> flowx;
	util::ndarray<float, 2> flowy;
	util::ndarray<math::vector<float, 2>, 2> velocity;
	util::ndarray<float, 2> sediment;
	util::ndarray<float, 2> new_sediment;

	gfx::painter painter;

	bool paused = true;
	bool show_water = true;
	bool show_velocity = false;
	bool show_particles = false;
	bool erosion_on = false;
	bool rain_on = false;

	struct particle
	{
		math::point<float, 2> position;
		int lifetime;
	};

	std::vector<particle> particles;
};

water_2d_app::water_2d_app(options const &, context const & ctx)
{
	(void)ctx;
	// ctx.vsync(true);

	simulation_area[0] = {-1.f, 1.f};
	simulation_area[1] = {-1.f, 1.f};

	inner_area = math::shrink(simulation_area, simulation_area.dimensions() / (2.f * N));

	bed.assign({N, N}, 0.f);
	water.assign({N, N}, 0.f);
	flowx.assign({N + 1, N}, 0.f);
	flowy.assign({N, N + 1}, 0.f);
	velocity.assign({N, N}, math::vector<float, 2>::zero());
	sediment.assign({N, N}, 0.f);
	new_sediment.assign({N, N}, 0.f);

	// water(N / 2, N / 2) = 100.f;

	random::uniform_sphere_vector_distribution<float, 2> d;
	util::ndarray<math::vector<float, 2>, 2> grads({9, 9});
	for (auto & v : grads)
		v = d(rng);

	pcg::perlin<float, 2> noise(std::move(grads), pcg::seamless);

	for (int y = 0; y < N; ++y)
	{
		for (int x = 0; x < N; ++x)
		{
			// if (x == N / 2 && (y != N / 4 && y != 3 * N / 4))
			// 	bed(x, y) = 100.f;

			[[maybe_unused]] float tx = (x + 0.5f) / N;
			[[maybe_unused]] float ty = (y + 0.5f) / N;

			float n = 0.f;

			// Archipelago
			n = 0.5f * noise(tx, ty) + 0.5f * noise(math::frac(2.f * tx), math::frac(2.f * ty));
			n -= 0.5f * math::distance(math::point{tx, ty}, math::point{0.5f, 0.5f});
			n = 5.f * math::smoothstep(math::clamp(math::unlerp({0.45f, 0.55f}, n), {0.f, 1.f}));

			// Island
			// n = 4.f * std::max(0.f, 1.f - 3.f * math::distance(math::point{tx, ty}, math::point{0.5f, 0.5f}) + (2.f * noise(tx, ty) - 1.f) * 0.25f);

			// Delta
			// n = 4.f * std::abs(2.f * noise(tx, ty) - 1.f) * (1.f - tx);

			// Double delta
			// n = 4.f * std::abs(2.f * noise(tx, ty) - 1.f) * std::pow(4.f * tx * (1.f - tx), 2.f);

			// River
			// n = std::min(4.f, math::lerp(10.f, 20.f, noise(tx, ty)) * std::abs(ty - math::lerp({0.4f, 0.6f}, noise(tx, 0.f))));

			// n = 0.f;

			// Canyon
			// n = 10.f * math::clamp(10.f * math::lerp(-1.f, 1.f, noise(tx, ty)), {0.f, 1.f});

			// Canyon v2
			// n = math::clamp(20.f * std::abs(2.f * noise(tx, ty) - 1.f) - 2.f, {0.f, 5.f});

			// Shore
			// n = math::clamp(20.f * (ty - math::lerp(0.4f, 0.6f, noise(0.f, tx))), {-4.f, 4.f}) + 4.f;

			// Circles
			// n = 20.f * std::max(0.f, 1.f - 25.f * math::length(math::pointwise_mult({1.f, 0.1f}, math::point{tx, std::fmod(10.f * ty, 1.f)} - math::point{0.5f, 0.5f})));

			bed(x, y) = n;

			// water(x, y) = std::max(0.f, 4.0f - bed(x, y));

			// water(x, y) = 10.f - bed(x, y);

			// if (x == 0)
			// 	water(x, y) = 100.f - bed(x, y);

			water(x, y) = std::max(0.f, 1.f - bed(x, y)) * (0.5f + 0.5f * noise(tx, ty));

			flowx(x, y) = (2.f * ty - 1.f) * water(x, y) * 10.f / N;
		}
	}

	// water(N / 2, N / 2) = 10000.f;
}

float const dt = 0.001f;

void water_2d_app::update()
{
	if (paused)
		return;

	float const dx = simulation_area[0].length() / N;
	float const dy = simulation_area[1].length() / N;
	float const g = 10.f;
	float const friction = std::pow(1.f, dt);
	float const viscosity = 0.f;
	float const sediment_capacity = 0.1f;
	float const erosion = 1.f;
	float const deposition = 10.f;
	float const coriolis = 0.01f;
	int const max_particles = 16 * 1024;
	int const particles_per_frame = 16;
	int const max_particle_lifetime = max_particles / particles_per_frame;

	// Init boundary flows
	for (int x = 0; x < N; ++x)
	{
		// flowy(x, 0) = -1.f / N;
		// flowy(x, N) =  1.f / N;

		// if (x >= 3 * N / 4) flowy(x, 0) = 3.f / N;
		// if (x < N / 4) flowy(x, N) = -3.f / N;

		// flowy(x, 0) = 5.f * std::sin(10.f * time) / N;

		flowy(x, 0) = 0.f;
		flowy(x, N) = 0.f;
	}

	for (int y = 0; y < N; ++y)
	{
		// flowx(0, y) =  5.f * std::sin(10.f * time) / N;

		// flowx(0, y) =  10.f / N;
		// flowx(N, y) =  1.f / N;

		// flowx(0, y) = 10.f * ((y + 0.5f) / N) / N;
		// flowx(N, y) = - (10.f / N - flowx(0, y));

		// if (y < N / 4) flowx(0, y) = 3.f / N;
		// if (y >= 3 * N / 4) flowx(N, y) = -3.f / N;

		// if (bed(0, y) < 0.75f)
		// 	flowx(0, y) = 4.f / N;
			// flowx(0, y) = 5.f * math::sqr(std::max(0.f, std::sin(15.f * time))) / N;

		// if (bed(N - 1, y) < 0.75f)
		// 	flowx(N, y) = 5.f / N;
	}

	// Rain :)
	if (rain_on)
		for (int y = 0; y < N; ++y)
			for (int x = 0; x < N; ++x)
				// if (random::uniform<float>(rng) < 0.1f)
					water(x, y) += random::uniform<float>(rng) * dt;
					// water(x, y) += dt;

	// water(N/2, N/2) += 1000.f * dt;

	// Update X flows
	for (int y = 0; y < N; ++y)
	{
		for (int x = 1; x < N; ++x)
			flowx(x, y) = friction * flowx(x, y) + g * dt * (water(x - 1, y) + bed(x - 1, y) - water(x, y) - bed(x, y));

		flowx(0, y) = friction * flowx(0, y) + g * dt * (water(N - 1, y) + bed(N - 1, y) - water(0, y) - bed(0, y));
	}

	// Update Y flows
	for (int y = 1; y < N; ++y)
		for (int x = 0; x < N; ++x)
			flowy(x, y) = friction * flowy(x, y) + g * dt * (water(x, y - 1) + bed(x, y - 1) - water(x, y) - bed(x, y));

	// Apply viscosity to X flows
	for (int y = 0; y < N; ++y)
	{
		for (int x = 1; x < N; ++x)
		{
			float h = (flowx(x, y) > 0.f) ? water(x - 1, y) : water(x, y);
			h *= h;

			if (h > 0.f)
				flowx(x, y) *= h / (h + 3.f * dt * viscosity);
		}
	}

	// Apply viscosity to Y flows
	for (int y = 1; y < N; ++y)
	{
		for (int x = 0; x < N; ++x)
		{
			float h = (flowy(x, y) > 0.f) ? water(x, y - 1) : water(x, y);
			h *= h;

			if (h > 0.f)
				flowy(x, y) *= h / (h + 3.f * dt * viscosity);
		}
	}

	// Scale flows
	for (int y = 0; y < N; ++y)
	{
		for (int x = 0; x < N; ++x)
		{
			float outflow = 0.f;
			outflow += std::max(0.f, - flowx(x    , y    ));
			outflow += std::max(0.f,   flowx((x + 1) % N, y    ));
			outflow += std::max(0.f, - flowy(x    , y    ));
			outflow += std::max(0.f,   flowy(x    , y + 1));

			float max_outflow = water(x, y) * dx * dy / dt;

			if (outflow > 0.f)
			{
				float scale = std::min(1.f, max_outflow / outflow);

				if (flowx(x, y) < 0.f) flowx(x, y) *= scale;
				if (flowx((x + 1) % N, y) > 0.f) flowx((x + 1) % N, y) *= scale;

				if (flowy(x, y) < 0.f) flowy(x, y) *= scale;
				if (flowy(x, y + 1) > 0.f) flowy(x, y + 1) *= scale;
			}
		}
	}

	// Update water and compute velocity
	for (int y = 0; y < N; ++y)
	{
		for (int x = 0; x < N; ++x)
		{
			float w_old = water(x, y);

			water(x, y) += dt / dx / dy * (flowx(x, y) + flowy(x, y) - flowx((x + 1) % N, y) - flowy(x, y + 1));

			float w_avg = (w_old + water(x, y)) / 2.f;

			math::vector v{0.f, 0.f};

			if (w_avg > 0.f)
			{
				v[0] = (flowx(x, y) + flowx((x + 1) % N, y)) / 2.f / dx / w_avg;
				v[1] = (flowy(x, y) + flowy(x, y + 1)) / 2.f / dy / w_avg;
			}

			velocity(x, y) = v;
		}
	}

	// Add coriolis force
	for (int y = 0; y < N; ++y)
	{
		for (int x = 0; x < N; ++x)
		{
			auto n = math::ort(velocity(x, y)) * coriolis * dt / 2.f * std::sin(float(math::pi) / 2.f * ((y + 0.5f) * 2.f / N - 1.f));

			flowx(x, y) += n[0];
			flowx((x + 1) % N, y) += n[0];
			flowy(x, y) += n[1];
			flowy(x, y + 1) += n[1];
		}
	}

	if (erosion_on)
	{
		// Erosion-deposition
		for (int y = 0; y < N; ++y)
		{
			for (int x = 0; x < N; ++x)
			{
				float c = sediment_capacity * water(x, y) * math::length(velocity(x, y));

				if (c >= sediment(x, y))
				{
					float delta = std::min(bed(x, y), dt * erosion * (c - sediment(x, y)));
					bed(x, y) -= delta;
					sediment(x, y) += delta;
				}
				else
				{
					float delta = std::min(sediment(x, y), dt * deposition * (sediment(x, y) - c));
					bed(x, y) += delta;
					sediment(x, y) -= delta;
				}
			}
		}

		// Sediment transport
		for (int y = 0; y < N; ++y)
		{
			for (int x = 0; x < N; ++x)
			{
				math::point p = math::lerp(simulation_area, {(x + 0.5f) / N, (y + 0.5f) / N});
				p -= velocity(x, y) * dt;

				p = math::clamp(p, inner_area);

				auto q = math::unlerp(inner_area, p) * (N - 1.f);

				int ix = std::min<int>(N - 2, std::floor(q[0]));
				int iy = std::min<int>(N - 2, std::floor(q[1]));

				float tx = q[0] - ix;
				float ty = q[1] - iy;

				new_sediment(x, y) = 0.f
					+ sediment(ix, iy) * (1.f - tx) * (1.f - ty)
					+ sediment(ix + 1, iy) * tx * (1.f - ty)
					+ sediment(ix, iy + 1) * (1.f - tx) * ty
					+ sediment(ix + 1, iy + 1) * tx * ty
					;
			}
		}

		std::swap(sediment, new_sediment);
	}

	// Init particles
	for (int i = 0; i < particles_per_frame; ++i)
		if (particles.size() < max_particles)
		{
			math::point pos{random::uniform<float>(rng, 0.f, N), random::uniform<float>(rng, 0.f, N)};
			if (water(std::floor(pos[0]), std::floor(pos[1])) > 1e-3f)
				particles.push_back({pos, 0});
		}

	// Update particles
	for (auto & p : particles)
	{
		p.position += velocity(std::min<int>(N - 1, std::floor(p.position[0])), std::min<int>(N - 1, std::floor(p.position[1]))) * (N * dt);
		p.position[0] = math::fmod(p.position[0], float(N));
		p.position[1] = math::clamp(p.position[1], {0.f, N});
		p.lifetime += 1;
	}

	// Destroy particles
	for (int i = 0; i < particles.size();)
	{
		if (particles[i].lifetime >= max_particle_lifetime)
		{
			std::swap(particles[i], particles.back());
			particles.pop_back();
		}
		else
			++i;
	}

	time += dt;
}

void water_2d_app::present()
{
	gl::ClearColor(0.8f, 0.8f, 0.8f, 0.8f);
	gl::Clear(gl::COLOR_BUFFER_BIT);

	math::box<float, 3> view_area;
	{
		view_area[0] = simulation_area[0];
		view_area[1] = simulation_area[1];
		view_area[2] = {-1.f, 1.f};

		float extra_y = view_area[1].length() * 0.f;
		view_area[1].min -= extra_y;
		view_area[1].max += extra_y;
		float extra_x = view_area[1].length() * aspect_ratio - view_area[0].length();
		view_area[0].min -= extra_x / 2.f;
		view_area[0].max += extra_x / 2.f;
	}

	math::matrix<float, 4, 4> const transform = math::orthographic{view_area}.homogeneous_matrix();

	painter.rect(simulation_area, {0, 0, 0, 255});

	float invN = 1.f / N;

	for (int y = 0; y < N; ++y)
	{
		for (int x = 0; x < N; ++x)
		{
			auto min = simulation_area.corner(x * invN, y * invN);
			auto max = simulation_area.corner((x + 1) * invN, (y + 1) * invN);

			float b = 1.f - std::exp(- 1.f * bed(x, y));
			float w = 1.f - std::exp(- 1.f * water(x, y));
			float s = 1.f - std::exp(- 1.f * sediment(x, y));

			// if (water(x, y) > 1e-3f)
			// 	w = 0.5f + 0.5f * w;

			auto bed_color = gfx::to_coloru8(gfx::color_4f{0.96f, 0.72f, 0.53f, b});
			// auto bed_color = gfx::to_coloru8(gfx::color_4f{0.3f, 0.5f, 0.1f, b});
			auto color = gfx::to_coloru8(gfx::color_4f{0.f, 0.25f, 1.f, w});
			auto sediment_color = gfx::to_coloru8(gfx::color_4f{1.f, 0.25f, 0.f, s * w});

			auto box = math::box<float, 2>::singleton(min) | math::box<float, 2>::singleton(max);

			painter.rect(box, bed_color);
			if (show_water)
			{
				painter.rect(box, color);
				painter.rect(box, sediment_color);
			}
		}
	}

	if (show_velocity)
	{
		for (int y = 0; y < N; ++y)
		{
			for (int x = 0; x < N; ++x)
			{
				auto center = simulation_area.corner((x + 0.5f) * invN, (y + 0.5f) * invN);

				auto v = water(x, y) * velocity(x, y) / 40.f;

				float max_length = 0.05f;
				auto l = math::length(v);

				// auto color = gfx::to_coloru8(gfx::color_4f{1.f, std::exp(- 10.f * l), 0.f, 1.f});
				auto color = gfx::color_rgba{255, 255, 255, 255};

				float s = std::min(1.f, l / max_length);
				// s = 1.f;

				auto d = math::normalized(v) * max_length * s / 2.f;

				painter.line(center - d, center + d, s * 0.001f, 0.f, color, color, false);
			}
		}
	}

	if (show_particles)
	{
		for (auto const & p : particles)
		{
			auto pos = math::lerp(simulation_area, math::vector{p.position[0] / N, p.position[1] / N});
			painter.circle(pos, 0.005f, {255, 255, 255, 127}, 6);
		}
	}

	auto mouse = math::cast<float>(state().mouse);

	mouse[0] = math::lerp(view_area[0], mouse[0] / state().size[0]);
	mouse[1] = math::lerp(view_area[1], 1.f - mouse[1] / state().size[1]);

	int mx = std::floor(N * math::unlerp(simulation_area[0], mouse[0]));
	int my = std::floor(N * math::unlerp(simulation_area[1], mouse[1]));

	if (mx >= 0 && mx < N && my >= 0 && my < N)
	{
		gfx::painter::text_options opts
		{
			.scale = {0.004f, -0.004f},
			.c = {255, 0, 255, 255},
		};
		painter.text({view_area[0].center(), math::lerp(view_area[1], 0.03f)}, std::format("b {:.3f} + w {:.3f} = h {:.3f}", bed(mx, my), water(mx, my), bed(mx, my) + water(mx, my)), opts);
		painter.text({view_area[0].center(), math::lerp(view_area[1], 0.05f)}, std::format("s {:.3f}", sediment(mx, my)), opts);

		if (state().mouse_button_down.contains(app::mouse_button::left))
		{
			for (int dy = -2; dy <= 2; ++dy)
			{
				for (int dx = -2; dx <= 2; ++dx)
				{
					int x = mx + dx;
					int y = my + dy;

					if (x >= 0 && x < N && y >= 0 && y < N)
						water(x, y) += 1000.f * dt;
				}
			}
		}
	}

	painter.render(transform);

	{
		painter.text({20.f, 20.f}, std::format("{} particles", particles.size()), {.scale = {2.f, 2.f}, .x = gfx::painter::x_align::left, .y = gfx::painter::y_align::top,  .c = {0, 0, 0, 255}});
		painter.render(math::window_camera{state().size[0], state().size[1]}.transform());
	}
}

void water_2d_app::on_event(app::resize_event const & event)
{
	app::application_base::on_event(event);

	aspect_ratio = (event.size[0] * 1.f) / event.size[1];
}

void water_2d_app::on_event(app::key_event const & event)
{
	if (event.down && event.key == app::keycode::SPACE)
		paused ^= true;

	if (event.down && event.key == app::keycode::W)
		show_water ^= true;

	if (event.down && event.key == app::keycode::V)
		show_velocity ^= true;

	if (event.down && event.key == app::keycode::P)
		show_particles ^= true;

	if (event.down && event.key == app::keycode::E)
		erosion_on ^= true;

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

namespace psemek::app
{

	std::unique_ptr<application::factory> make_application_factory()
	{
		return default_application_factory<water_2d_app>({.name = "Water 2D example", .multisampling = 4});
	}

}