#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/gradient.hpp>
#include <psemek/random/generator.hpp>
#include <psemek/random/device.hpp>
#include <psemek/random/uniform_ball.hpp>
#include <psemek/pcg/perlin.hpp>
#include <psemek/pcg/fractal.hpp>
#include <psemek/util/ndarray.hpp>
#include <psemek/log/log.hpp>
#include <psemek/io/file_stream.hpp>

using namespace psemek;

auto make_perlin(random::generator & rng, int min_octave, int max_octave, float power)
{
	std::vector<pcg::perlin<float, 2>> octaves;
	std::vector<float> weights;

	random::uniform_sphere_vector_distribution<float, 2> random_vector{};

	for (int octave = min_octave; octave < max_octave; ++octave)
	{
		int size = 1 << octave;
		util::ndarray<math::vector<float, 2>, 2> gradients({size + 1, size + 1});
		for (auto & g : gradients)
			g = random_vector(rng);
		octaves.emplace_back(std::move(gradients));
		weights.push_back(std::pow(power, - octave));
	}

	float weight_sum = 0.f;
	for (auto w : weights)
		weight_sum += w;
	for (auto & w : weights)
		w /= weight_sum;

	return pcg::fractal<pcg::perlin<float, 2>>(std::move(octaves), std::move(weights));
}

void make_force_field(random::generator & rng, util::ndarray<math::vector<float, 2>, 2> & result, float scale)
{
	auto noise_1 = make_perlin(rng, 4, 6, 2.f);
	auto noise_2 = make_perlin(rng, 4, 6, 2.f);

	for (int y = 0; y < result.height(); ++y)
	{
		for (int x = 0; x < result.width(); ++x)
		{
			math::point p{(x + 0.5f) / result.height(), (y + 0.5f) / result.width()};
			result(x, y) = scale * (math::vector{noise_1(p), noise_2(p)} * 2.f - math::vector{1.f, 1.f});
		}
	}
}

struct weather_app
	: app::application_base
{
	static constexpr int N = 128;

	const bool static_mode = true;

	const float dt = 20.f;
	const float viscosity = static_mode ? 0.005f : 0.f;
	const float advection_magnification = 1.f;
	const float temperature_diffusion = 0.001f;
	const float cooling = 0.01f / 300.f;
	const float cooling_factor = std::exp(- cooling * dt);
	const float heating = 323.f * (std::exp(cooling * dt) - 1.f) / dt;
	const float coriolis = 0.f;
	const float coriolis_bands = 2.f;
	const float friction = 0.f;
	const float slope_force = 0.05f;
	const float vorticity_confinement = 0.f;
	const float elevation_temperature_drop = 30.f;
	const float evaporation = 1.0f;
	const float max_humidity_factor = 100.f;
	const float precipitation_factor = 0.0001f;
	const float force_field_amplitude = 0.00005f;
	const float random_forces = 0.25f * (static_mode ? 0.f : 1.f);
	const float force_field_switch_duration = 720.f * 7.5f; // 7.5 days
	const int force_field_switch_frames = std::round(force_field_switch_duration / dt);
	// const float friction_factor = 1.f - std::exp(- friction * dt);

	const bool periodic_x = true;

	// random::generator rng{random::device{}};
	random::generator rng{0, 0};

	gfx::pixmap_rgba biomes_map;

	float expected_temperature_at(int y) const
	{
		// float latitude = (y - N * 0.5f) * 2.f / N;
		// return std::cos(latitude * float(math::pi));

		return temperature_income_at(y) * dt / (std::exp(cooling * dt) - 1.f);
	}

	float temperature_income_at(int y) const
	{
		// float latitude = (y - N * 0.5f) * 2.f / N;
		float latitude = y * 1.f / N;
		// return heating * math::lerp(0.75f, 1.f, std::cos(latitude * float(math::pi) / 2.f));
		return heating * math::lerp(0.8f, 1.f, 1.f - std::abs(latitude));
	}

	int wrap(int i) const
	{
		return (i + N) % N;
	}

	weather_app(options const &, context const &)
	{
		simulation_box_ = {{{0.f, N}, {0.f, N}}};

		terrain_.resize({N, N}, 0.f);
		velocity_.resize({N, N});
		new_velocity_.resize({N, N});
		pressure_.resize({N, N}, 0.f);
		temperature_.resize({N, N}, -1.f);
		new_temperature_.resize({N, N});
		average_temperature_.resize({N, N}, 0.f);
		force_field_main_.resize({N, N});
		force_field_current_.resize({N, N});
		force_field_next_.resize({N, N});
		vorticity_.resize({N, N});
		humidity_.resize({N, N});
		new_humidity_.resize({N, N});
		precipitation_.resize({N, N});
		average_precipitation_.resize({N, N});

		// random::generator rng{0, 0};
		random::uniform_ball_vector_distribution<float, 2> random_velocity{};

		// for (auto & v : velocity_)
		// {
		// 	v = random_velocity(rng) * 0.01f;
		// 	v += std::cos(0.5f * float(math::pi) * latitude * coriolis_bands
		// }

		for (int y = 0; y < N; ++y)
		{
			for (int x = 0; x < N; ++x)
			{
				float latitude = (N * 0.5f - y) * 2.f / N;
				velocity_(x, y) = random_velocity(rng) * 0.f + 0.f * math::vector{-std::cos(0.5f * float(math::pi) * latitude * coriolis_bands), 0.f};
				temperature_(x, y) = expected_temperature_at(y);
			}
		}

		auto terrain_noise = make_perlin(rng, 2, 10, 1.6f);

		for (int y = 0; y < N; ++y)
		{
			for (int x = 0; x < N; ++x)
			{
				auto d = math::length(math::vector{x - N / 2.f, y - N / 2.f}) / (N / 2.f);
				(void)d;
				float value = terrain_noise((x + 0.5f) / N, (y + 0.5f) / N);
				value = pow(value, 4.f) - d / 4.f;
				value = math::lerp(1.f, 16.f, value);
				terrain_(x, y) = value;
			}
		}

		auto heightmap = gfx::read_image<std::uint8_t>(io::file_istream{std::filesystem::path{PSEMEK_EXAMPLES_DIR} / "heightmap_seed_3.png"});

		for (int y = 0; y < N; ++y)
		{
			for (int x = 0; x < N; ++x)
			{
				terrain_(x, y) = ((heightmap(x, y) / 255.f) * 2048.f - 512.f) / 1024.f;
			}
		}

		for (int y = 0; y < N; ++y)
		{
			for (int x = 0; x < N; ++x)
			{
				if (terrain_(x, y) > 0.f)
					continue;

				float max_humidity = std::max(0.f, temperature_(x, y) - 223.f) * max_humidity_factor;
				humidity_(x, y) = max_humidity + dt * evaporation * std::max(0.f, temperature_(x, y) - 273.f) * (1.f - precipitation_factor * dt) / precipitation_factor / dt;
			}
		}

		make_force_field(rng, force_field_main_, 0.5f);
		make_force_field(rng, force_field_next_, 0.5f);

		biomes_map = gfx::read_image<gfx::color_rgba>(io::file_istream{std::filesystem::path{PSEMEK_EXAMPLES_DIR} / "biomes.png"});
	}

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

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

		if (event.down && event.key == app::keycode::T)
			show_temperature_ ^= true;

		if (event.down && event.key == app::keycode::D)
			show_temperature_delta_ ^= true;

		if (event.down && event.key == app::keycode::A)
			show_average_temperature_delta_ ^= true;

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

		if (event.down && event.key == app::keycode::H)
			show_land_ ^= true;

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

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

		if (event.down && event.key == app::keycode::Q)
			show_average_precipitation_ ^= true;

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

	void update() override
	{
		if (paused_)
			return;

		// Update force field
		if ((frame_ % force_field_switch_frames) == 0)
		{
			std::swap(force_field_current_, force_field_next_);
			make_force_field(rng, force_field_next_, 0.5f);
		}
		float const force_field_t = ((frame_ % force_field_switch_frames) + 0.5f) / force_field_switch_frames;

		int xmin = periodic_x ? 0 : 1;
		int xmax = periodic_x ? N : N - 1; // exclusive

		// Temperature source
		for (int y = 0; y < N; ++y)
		{
			for (int x = 0; x < N; ++x)
			{
				// temperature_(x, y) = math::lerp(temperature_(x, y), expected_temperature_at(y), 1.f - std::exp(- heating * dt));
				temperature_(x, y) += dt * temperature_income_at(y);
				temperature_(x, y) *= cooling_factor;
			}
		}

		// Evaporation
		for (int y = 0; y < N; ++y)
		{
			for (int x = 0; x < N; ++x)
			{
				if (terrain_(x, y) <= 0.f)
					humidity_(x, y) += dt * evaporation * std::max(0.f, temperature_(x, y) - 273.f);

				// float discharge = std::min(humidity_(x, y), precipitation_factor * dt);
				// float discharge = humidity_(x, y) * precipitation_factor * dt;
				float max_humidity = std::max(0.f, temperature_(x, y) - 223.f) * max_humidity_factor;
				float discharge = std::max(0.f, humidity_(x, y) - max_humidity) * precipitation_factor * dt;
				humidity_(x, y) -= discharge;
				precipitation_(x, y) = discharge / dt;
			}
		}

		// Velocity & temperature advection
		for (int i = 0; i < N; ++i)
		{
			new_temperature_(i, 0) = temperature_(i, 0);
			new_temperature_(i, N - 1) = temperature_(i,  N - 1);

			new_humidity_(i, 0) = humidity_(i, 0);
			new_humidity_(i, N - 1) = humidity_(i,  N - 1);

			if (!periodic_x)
			{
				new_temperature_(0, i) = temperature_(0, i);
				new_temperature_(N - 1, i) = temperature_(N - 1, i);

				new_humidity_(0, i) = humidity_(0, i);
				new_humidity_(N - 1, i) = humidity_(N - 1, i);
			}
		}
		for (int y = 1; y < N - 1; ++y)
		{
			for (int x = xmin; x < xmax; ++x)
			{
				auto v = velocity_(x, y);
				auto p = math::point{x + 0.5f, y + 0.5f} - (advection_magnification * dt) * v;
				p[0] = p[0] - 0.5f;
				p[1] = math::clamp(p[1] - 0.5f, {0.f, N - 1.f});

				if (!periodic_x)
					p[0] = math::clamp(p[0], {0.f, N - 1.f});

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

				if (!periodic_x)
					ix = std::min(N - 1, ix);

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

				new_velocity_(x, y) = math::lerp(
					math::lerp(velocity_(wrap(ix + 0), iy + 0), velocity_(wrap(ix + 1), iy + 0), tx),
					math::lerp(velocity_(wrap(ix + 0), iy + 1), velocity_(wrap(ix + 1), iy + 1), tx),
					ty
				);

				new_temperature_(x, y) = math::lerp(
					math::lerp(temperature_(wrap(ix + 0), iy + 0), temperature_(wrap(ix + 1), iy + 0), tx),
					math::lerp(temperature_(wrap(ix + 0), iy + 1), temperature_(wrap(ix + 1), iy + 1), tx),
					ty
				);

				new_humidity_(x, y) = math::lerp(
					math::lerp(humidity_(wrap(ix + 0), iy + 0), humidity_(wrap(ix + 1), iy + 0), tx),
					math::lerp(humidity_(wrap(ix + 0), iy + 1), humidity_(wrap(ix + 1), iy + 1), tx),
					ty
				);
			}
		}
		std::swap(velocity_, new_velocity_);
		std::swap(temperature_, new_temperature_);
		std::swap(humidity_, new_humidity_);

		// Apply velocity diffusion
		for (int y = 0; y < N; ++y)
			for (int x = 0; x < N; ++x)
				new_velocity_(x, y) = velocity_(x, y);
		for (int y = 1; y < N - 1; ++y)
		{
			for (int x = xmin; x < xmax; ++x)
			{
				// Velocity Laplacian
				auto laplacian = velocity_(wrap(x + 1), y) + velocity_(wrap(x - 1), y) + velocity_(x, y + 1) + velocity_(x, y - 1) - 4.f * velocity_(x, y);
				new_velocity_(x, y) = velocity_(x, y) + viscosity * dt * laplacian;
			}
		}
		std::swap(velocity_, new_velocity_);

		// Apply temperature diffusion
		for (int y = 0; y < N; ++y)
			for (int x = 0; x < N; ++x)
				new_temperature_(x, y) = temperature_(x, y);
		for (int y = 1; y < N - 1; ++y)
		{
			for (int x = xmin; x < xmax; ++x)
			{
				// Temperature Laplacian
				auto laplacian = temperature_(wrap(x + 1), y) + temperature_(wrap(x - 1), y) + temperature_(x, y + 1) + temperature_(x, y - 1) - 4.f * temperature_(x, y);
				new_temperature_(x, y) = temperature_(x, y) + temperature_diffusion * dt * laplacian;
			}
		}
		std::swap(temperature_, new_temperature_);

		// Compute vorticity
		for (int y = 1; y < N - 1; ++y)
			for (int x = xmin; x < xmax; ++x)
				vorticity_(x, y) = (velocity_(x, y + 1)[0] - velocity_(x, y - 1)[0]) / 2.f - (velocity_(wrap(x + 1), y)[1] - velocity_(wrap(x - 1), y)[1]) / 2;

		// Apply forces & friction
		for (int y = 1; y < N - 1; ++y)
		{
			for (int x = xmin; x < xmax; ++x)
			{
				float latitude = (N * 0.5f - y) * 2.f / N;
				// float latitude = (N - y) * 1.f / N;
				// velocity_(x, y) += math::ort(velocity_(x, y)) * (coriolis * dt * std::sin(0.5f * float(math::pi) * latitude * coriolis_bands));
				velocity_(x, y) = math::rotate(velocity_(x, y), coriolis * dt * std::sin(0.5f * float(math::pi) * latitude * coriolis_bands));

				auto force = force_field_main_(x, y) + random_forces * math::lerp(force_field_current_(x, y), force_field_next_(x, y), force_field_t);
				velocity_(x, y) += (dt * force_field_amplitude) * force;

				math::vector terrain_gradient
				{
					(std::max(0.f, terrain_(x + 1, y)) - std::max(0.f, terrain_(x - 1, y))) / 2.f,
					(std::max(0.f, terrain_(x, y + 1)) - std::max(0.f, terrain_(x, y - 1))) / 2.f,
				};

				// [[maybe_unused]] float slope_factor = std::exp(- dt * slope_force * math::dot(math::normalized(velocity_(x, y)), terrain_gradient));
				// velocity_(x, y) *= std::min(1.f, slope_factor);
				// velocity_(x, y) -= terrain_gradient * slope_force * dt;

				// [[maybe_unused]] float slope_factor = std::exp(- dt * slope_force * std::pow(math::length(terrain_gradient), 4.f));
				[[maybe_unused]] float slope_factor = std::exp(- dt * slope_force * std::pow(std::max(0.f, terrain_(x, y)), 1.f));
				velocity_(x, y) *= slope_factor;

				// Directional external force
				// velocity_(x, y)[1] += 0.001f * dt * std::sin(0.5f * float(math::pi) * latitude * coriolis_bands);
				// velocity_(x, y) += math::direction(frame_ * dt * 2.f * float(math::pi) / 10080.f) * 0.001f * dt;

				math::vector vorticity_gradient
				{
					(vorticity_(wrap(x + 1), y) - vorticity_(wrap(x - 1), y)) / 2.f,
					(vorticity_(x, y + 1) - vorticity_(x, y - 1)) / 2.f,
				};
				if (auto l = math::length(vorticity_gradient); l > 0.f)
					vorticity_gradient /= l;

				velocity_(x, y) += vorticity_confinement * dt * vorticity_(x, y) * math::ort(vorticity_gradient);

				float local_friction = friction * terrain_(x, y);
				// velocity_(x, y) -= local_friction * velocity_(x, y) * math::length(velocity_(x, y));
				float local_friction_factor = std::exp(- local_friction * dt);
				velocity_(x, y) *= local_friction_factor;
			}
		}

		// Solve Poisson equation for pressure
		for (int iteration = 0; iteration < 16; ++iteration)
		{
			int ymin = ((iteration % 2) == 0) ? 1 : N - 2;
			int ymax = ((iteration % 2) == 0) ? N - 1 : 0;
			int ystep = ((iteration % 2) == 0) ? 1 : -1;

			for (int y = ymin; y != ymax; y += ystep)
			{
				for (int x = xmin; x < xmax; ++x)
				{
					// Velocity divergence
					float divergence = (velocity_(wrap(x + 1), y)[0] - velocity_(wrap(x - 1), y)[0] + velocity_(x, y + 1)[1] - velocity_(x, y - 1)[1]) / 2.f;

					// Gauss-Seidel iteration step
					pressure_(x, y) = (pressure_(wrap(x - 1), y) + pressure_(wrap(x + 1), y) + pressure_(x, y - 1) + pressure_(x, y + 1) - divergence) / 4.f;
				}
			}
		}

		// Apply boundary conditions for pressure
		for (int i = 0; i < N; ++i)
		{
			if (!periodic_x)
			{
				pressure_(0, i) = pressure_(1, i);
				pressure_(N - 1, i) = pressure_(N - 2, i);
			}
			pressure_(i, 0) = pressure_(i, 1);
			pressure_(i, N - 1) = pressure_(i, N - 2);
		}
		if (!periodic_x)
		{
			pressure_(0, 0) = (pressure_(0, 1) + pressure_(1, 0)) / 2.f;
			pressure_(N-1, 0) = (pressure_(N-1, 1) + pressure_(N-2, 0)) / 2.f;
			pressure_(0, N-1) = (pressure_(0, N-2) + pressure_(1, N-2)) / 2.f;
			pressure_(N-1, N-1) = (pressure_(N-1, N-2) + pressure_(N-2, N-1)) / 2.f;
		}

		// Normalize pressure
		float average_pressure = 0.f;
		for (auto const & value : pressure_)
			average_pressure += value;
		average_pressure /= (1.f * N * N);
		for (auto & value : pressure_)
			value -= average_pressure;

		// Project velocity into divergence-free space
		// by subtracting pressure gradient
		for (int y = 1; y < N - 1; ++y)
		{
			for (int x = xmin; x < xmax; ++x)
			{
				// Pressure gradient
				math::vector gradient{
					(pressure_(wrap(x + 1), y) - pressure_(wrap(x - 1), y)) / 2.f,
					(pressure_(x, y + 1) - pressure_(x, y - 1)) / 2.f
				};
				velocity_(x, y) -= gradient;
			}
		}

		// Apply boundary conditions for velocity
		for (int i = 0; i < N; ++i)
		{
			if (!periodic_x)
			{
				float left_boundary_flow = 0.01f;//0.01f * std::sin((i * 1.f / N) * float(math::pi) * 4.f);
				float right_boundary_flow = -left_boundary_flow;

				velocity_(1, i)[0] = left_boundary_flow;
				velocity_(N-2, i)[0] = right_boundary_flow;

				velocity_(0, i)[0] = - velocity_(1, i)[0];
				velocity_(N-1, i)[0] = - velocity_(N-2, i)[0];
			}
			velocity_(i, 0)[1] = -velocity_(i, 1)[1];
			velocity_(i, N-2)[1] = -velocity_(i, N-2)[1];

			velocity_(i, 1)[0] = 0.01f;
			velocity_(i, N-2)[0] = 0.01f;
		}

		// Uncomment to visualize the force field
		// for (int y = 0; y < N; ++y)
		// 	for (int x = 0; x < N; ++x)
		// 		velocity_(x, y) = 100000.f * force_field_(x, y);

		// Uncomment to visualize the terrain gradient field
		// for (int y = 1; y < N - 1; ++y)
		// {
		// 	for (int x = 1; x < N - 1; ++x)
		// 	{
		// 		math::vector terrain_gradient
		// 		{
		// 			(terrain_(x + 1, y) - terrain_(x - 1, y)) / 2.f,
		// 			(terrain_(x, y + 1) - terrain_(x, y - 1)) / 2.f,
		// 		};

		// 		velocity_(x, y) = terrain_gradient;
		// 	}
		// }

		++frame_;

		// Update all-time average temperature & precipitation
		for (int y = 0; y < N; ++y)
		{
			for (int x = 0; x < N; ++x)
			{
				// float t = 1.f / frame_;
				// float t = 1.f / std::min(8192, frame_);
				float t = 1.f / std::min<float>(86400.f / dt, frame_); // year average

				if (static_mode)
					t = 1.f;

				average_temperature_(x, y) = math::lerp(average_temperature_(x, y), temperature_(x, y), t);
				average_precipitation_(x, y) = math::lerp(average_precipitation_(x, y), precipitation_(x, y), t);
			}
		}
	}

	void present() override
	{
		gl::ClearColor(0.f, 0.f, 0.f, 0.f);
		gl::Clear(gl::COLOR_BUFFER_BIT);

		float const aspect_ratio = state().size[0] * 1.f / state().size[1];
		math::box<float, 2> view_box = math::expand(simulation_box_, 1.f);
		if (view_box[0].length() / view_box[1].length() > aspect_ratio)
			view_box[1] = math::expand(view_box[1], (view_box[0].length() / aspect_ratio - view_box[1].length()) / 2.f);
		else
			view_box[0] = math::expand(view_box[0], (view_box[1].length() * aspect_ratio - view_box[0].length()) / 2.f);

		std::optional<math::vector<int, 2>> mouseover_cell;
		{
			auto mouse = math::lerp(view_box, math::vector{state().mouse[0] * 1.f / state().size[0], 1.f - state().mouse[1] * 1.f / state().size[1]});
			int x = std::floor(mouse[0]);
			int y = std::floor(mouse[1]);

			if (x >= 0 && x < N && y >= 0 && y < N)
				mouseover_cell = {x, y};
		}

		if (mouseover_cell && (state().mouse_button_down.contains(app::mouse_button::left) || state().mouse_button_down.contains(app::mouse_button::right)))
		{
			float delta = state().mouse_button_down.contains(app::mouse_button::left) ? 1.f : -1.f;

			int R = 4;

			for (int iy = -R; iy <= R; ++iy)
			{
				for (int ix = -R; ix <= R; ++ix)
				{
					int x = (*mouseover_cell)[0] + ix;
					int y = (*mouseover_cell)[1] + iy;

					if (x >= 0 && x < N && y >= 0 && y < N)
					{
						auto d = math::vector<float, 2>{ix, iy} / 2.f;
						terrain_(x, y) += delta * 0.05f * std::exp(- math::dot(d, d));
					}
				}
			}
		}

		if (mouseover_cell && state().mouse_button_down.contains(app::mouse_button::middle))
		{
			int R = 4;

			float average = 0.f;
			int count = 0;

			for (int iy = -R; iy <= R; ++iy)
			{
				for (int ix = -R; ix <= R; ++ix)
				{
					int x = (*mouseover_cell)[0] + ix;
					int y = (*mouseover_cell)[1] + iy;

					if (x >= 0 && x < N && y >= 0 && y < N)
					{
						average += terrain_(x, y);
						count += 1;
					}
				}
			}

			average /= count;

			for (int iy = -R; iy <= R; ++iy)
			{
				for (int ix = -R; ix <= R; ++ix)
				{
					int x = (*mouseover_cell)[0] + ix;
					int y = (*mouseover_cell)[1] + iy;

					if (x >= 0 && x < N && y >= 0 && y < N)
					{
						auto d = math::vector<float, 2>{ix, iy} / 2.f;
						terrain_(x, y) += (average - terrain_(x, y)) * 0.05f * std::exp(- math::dot(d, d));
					}
				}
			}
		}

		[[maybe_unused]] float const pixel_size = view_box[0].length() / state().size[0];

		auto map_color = [](float value, gfx::color_4f const & negative, gfx::color_4f const & positive){
			return math::lerp(negative, positive, 1.f/ (1.f + std::exp(- value)));
		};

		auto map_temperature = [&](float value) {
			return map_color(2.f * std::round(value / 20.f), {0.125f, 0.5f, 1.f, 0.75f}, {1.f, 0.5f, 0.125f, 0.75f});
		};

		auto map_biome = [this](float temperature, float precipitation)
		{
			auto x = math::clamp<int>(math::unlerp({  -3.f,  3.f}, precipitation) * biomes_map.width() , {0, biomes_map.width()  - 1});
			auto y = math::clamp<int>(math::unlerp({-10.f, 30.f}, temperature  ) * biomes_map.height(), {0, biomes_map.height() - 1});

			return gfx::to_colorf(biomes_map(x, y));
		};

		for (int y = 0; y < N; ++y)
		{
			for (int x = 0; x < N; ++x)
			{
				gfx::color_4f color = gfx::color_4f::zero();

				if (show_land_ || show_biomes_)
				{
					if (!show_biomes_)
					{
						if (terrain_(x, y) <= 0.f)
							color = {0.5f, 0.5f, 1.f, 1.f};
						else
							color = {1.f, 1.f, 1.f, 1.f};
					}
					else
					{
						if (terrain_(x, y) <= 0.f)
							color = map_color(8.f * terrain_(x, y), {0.f, 0.f, 0.125f, 1.f}, {0.f, 1.f, 1.5f, 1.f});
						else
						{
							float temperature = average_temperature_(x, y) - 273.f - std::max(0.f, terrain_(x, y)) * elevation_temperature_drop;
							// float precipitation = (std::log10(std::max(1e-9f, average_precipitation_(x, y))) + 3.f) * 2.f;
							// float precipitation = std::pow(average_precipitation_(x, y) * 125.f, 2.f) * 8.f;
							float precipitation = std::log2(std::max(1e-9f, average_precipitation_(x, y)));
							color = map_biome(temperature, precipitation);
						}
					}

					if (show_land_ && x > 0 && x + 1 < N && y > 0 && y + 1 < N)
					{
						math::vector terrain_gradient
						{
							(terrain_(x + 1, y) - terrain_(x - 1, y)) / 2.f,
							(terrain_(x, y + 1) - terrain_(x, y - 1)) / 2.f,
						};

						auto terrain_normal = math::normalized(math::vector{-terrain_gradient[0], -terrain_gradient[1], 0.125f});

						float lightness = 0.5f + 0.5f * math::dot(terrain_normal, math::normalized(math::vector{1.f, 2.f, 3.f}));

						color = gfx::dark(color, 1.f - lightness);
					}
				}

				if (show_temperature_)
					color = map_temperature(temperature_(x, y) - 273.f);

				if (show_temperature_delta_)
					color = gfx::blend(color, map_color((temperature_(x, y) - expected_temperature_at(y)), {0.125f, 0.5f, 1.f, 0.75f}, {1.f, 0.5f, 0.125f, 0.75f}));

				if (show_average_temperature_delta_)
					color = gfx::blend(color, map_color((average_temperature_(x, y) - expected_temperature_at(y)), {0.125f, 0.5f, 1.f, 0.75f}, {1.f, 0.5f, 0.125f, 0.75f}));

				if (show_pressure_)
					color = gfx::blend(color, map_color(10000.f * pressure_(x, y), {0.f, 0.f, 1.f, 0.75f}, {1.f, 0.f, 0.f, 0.75f}));

				if (show_water_vapor_)
					color = gfx::blend(color, map_color(humidity_(x, y) * 0.00001f, {0.f, 1.f, 1.f, -0.75f}, {0.f, 1.f, 1.f, 0.75f}));

				if (show_precipitation_)
				{
					// color = gfx::blend(color, map_color(precipitation_(x, y), {1.f, 1.f, 1.f, -1.f}, {1.f, 1.f, 1.f, 1.f}));

					float alpha = 2.f / (1.f + std::exp(- 0.1f * precipitation_(x, y))) - 1.f;

					gfx::color_4f cloud_color{1.f, 1.f, 1.f, alpha * 0.875f};

					if (y > 0 && y + 1 < N)
					{
						if (periodic_x || (x > 0 && x + 1 < N))
						{
							math::vector gradient
							{
								(precipitation_(wrap(x + 1), y) - precipitation_(wrap(x - 1), y)) / 2.f,
								(precipitation_(x, y + 1) - precipitation_(x, y - 1)) / 2.f,
							};

							auto normal = math::normalized(math::vector{-gradient[0], -gradient[1], 2.f});
							auto lightness = 0.5f + 0.5f * math::dot(normal, math::normalized(math::vector{1.f, 2.f, 3.f}));
							cloud_color = gfx::dark(cloud_color, 1.f - lightness);
						}
					}

					color = gfx::blend(color, cloud_color);
				}

				if (show_average_precipitation_)
					color = gfx::blend(color, map_color(average_precipitation_(x, y), {0.f, 1.f, 1.f, -0.75f}, {0.f, 1.f, 1.f, 0.75f}));

				painter_.rect({{{x, x + 1.f}, {y, y + 1.f}}}, gfx::to_coloru8(color));
			}
		}

		if (show_velocity_)
		{
			for (int y = 0; y < N; ++y)
			{
				for (int x = 0; x < N; ++x)
				{
					math::point center{x + 0.5f, y + 0.5f};
					auto v = velocity_(x, y);
					auto color = gfx::color_4f::zero();
					if (auto l = math::length(v); l > 0.f)
					{
						float const magnification = 1000.f;
						float const max_length = 1.5f;
						v *= 0.5f * max_length * (1.f - std::exp(- magnification * l)) / l;

						// color = gfx::lerp(gfx::color_4f{0.5f, 1.f, 0.f, 1.f}, gfx::color_4f{1.f, 0.f, 0.f, 1.f}, 1.f - std::exp(- 0.25f * magnification * l));
						color = map_color(0.1f * (temperature_(x, y) - 273.f), {0.125f, 0.5f, 1.f, 1.f}, {1.f, 0.5f, 0.125f, 1.f});
					}
					auto n = math::ort(v) * 0.3f;

					painter_.triangle(center - v - n, center - v + n, center + v, gfx::to_coloru8(color));
				}
			}
		}

		auto push_text = [&, row = 0](std::string const & text) mutable
		{
			painter_.text(view_box.corner(0.f, 1.f) - math::vector{0.f, row * pixel_size * 2.f * 12.f}, text, {.scale = {2.f * pixel_size, - 2.f * pixel_size}, .x = gfx::painter::x_align::left, .y = gfx::painter::y_align::top, .c = {255, 255, 255, 255}});
			++row;
		};

		push_text(std::format("Frame {}", frame_));
		push_text(std::format("Day {:.2f}", frame_ / (720.f / dt)));

		if (mouseover_cell)
		{
			int x = (*mouseover_cell)[0];
			int y = (*mouseover_cell)[1];
			painter_.rect({{{x, x + 1.f}, {y, y + 1.f}}}, {255, 255, 255, 127});

			push_text(std::format("{} {}", x, y));
			push_text(std::format("V = {:.3f} {:.3f}", velocity_(x, y)[0] * 1000.f, velocity_(x, y)[1] * 1000.f));
			push_text(std::format("P = {:.3f}", pressure_(x, y) * 1000.f));
			push_text(std::format("T = {:.3f}", temperature_(x, y) - 273.f));
			push_text(std::format("A = {:.3f}", average_temperature_(x, y) - 273.f));
			push_text(std::format("E = {:.3f}", expected_temperature_at(y) - 273.f));
			push_text(std::format("H = {:.3f}", terrain_(x, y)));
			push_text(std::format("W = {:.3f}", humidity_(x, y)));
			push_text(std::format("R = {:.3f}", precipitation_(x, y)));
			push_text(std::format("AR= {:.3f}", average_precipitation_(x, y)));
		}

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

private:
	gfx::painter painter_;

	math::box<float, 2> simulation_box_;

	bool paused_ = true;
	bool show_velocity_ = true;
	bool show_temperature_ = false;
	bool show_temperature_delta_ = false;
	bool show_average_temperature_delta_ = false;
	bool show_pressure_ = false;
	bool show_land_ = true;
	bool show_biomes_ = true;
	bool show_water_vapor_ = false;
	bool show_precipitation_ = false;
	bool show_average_precipitation_ = false;

	util::ndarray<float, 2> terrain_;
	util::ndarray<math::vector<float, 2>, 2> velocity_;
	util::ndarray<math::vector<float, 2>, 2> new_velocity_;
	util::ndarray<float, 2> pressure_;
	util::ndarray<float, 2> vorticity_;

	util::ndarray<float, 2> temperature_;
	util::ndarray<float, 2> new_temperature_;
	util::ndarray<float, 2> average_temperature_;

	util::ndarray<float, 2> humidity_;
	util::ndarray<float, 2> new_humidity_;
	util::ndarray<float, 2> precipitation_;
	util::ndarray<float, 2> average_precipitation_;

	util::ndarray<math::vector<float, 2>, 2> force_field_main_;
	util::ndarray<math::vector<float, 2>, 2> force_field_current_;
	util::ndarray<math::vector<float, 2>, 2> force_field_next_;

	int frame_ = 0;
};

namespace psemek::app
{

	std::unique_ptr<application::factory> make_application_factory()
	{
		return default_application_factory<weather_app>({.name = "Weather simulation test"});
	}

}