#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/random/generator.hpp>
#include <psemek/random/device.hpp>
#include <psemek/random/uniform_ball.hpp>
#include <psemek/util/ndarray.hpp>
#include <psemek/log/log.hpp>

using namespace psemek;

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

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

		obstacle_.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});

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

		for (auto & v : velocity_)
			v = random_velocity(rng);

		for (int y = 0; y < N; ++y)
		{
			for (int x = 0; x < N; ++x)
			{
				// velocity_(x, y)[0] += (y*1.f/N-0.5f);
				temperature_(x, y) = 0.5f - std::abs(y*2.f/N - 1.f);
			}
		}
	}

	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::P)
			show_pressure_ ^= true;
	}

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

		float const dt = 0.2f;
		float const viscosity = 0.f;
		float const temperature_diffusion = 0.01f;
		float const heating = 0.001f;
		float const coriolis = 0.1f;
		float const coriolis_bands = 5.f;
		float const friction = 0.1f;
		float const friction_factor = 1.f - std::exp(- friction * dt);

		// Temperature source
		for (int y = 0; y < N; ++y)
		{
			for (int x = 0; x < N; ++x)
			{
				temperature_(x, y) += dt * heating * (0.5f - std::abs(y*2.f/N - 1.f));
			}
		}

		auto wrap = [](int i)
		{
			return (i + N) % N;
		};

		// Velocity & temperature advection
		for (int y = 1; y < N - 1; ++y)
		{
			for (int x = 0; x < N; ++x)
			{
				auto v = velocity_(x, y);
				auto p = math::point{x + 0.5f, y + 0.5f} - dt * v;
				p[0] = math::clamp(p[0] - 0.5f, {0.f, N - 1.f});
				p[1] = math::clamp(p[1] - 0.5f, {0.f, N - 1.f});

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

				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
				);
			}
		}
		std::swap(velocity_, new_velocity_);
		std::swap(temperature_, new_temperature_);

		// 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 = 0; x < N; ++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 = 0; x < N; ++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_);

		// Apply forces & friction
		for (int y = 1; y < N - 1; ++y)
		{
			for (int x = 0; x < N; ++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));

				bool is_land = math::length(math::vector{x - N / 2.f, y - N / 4.f}) < N / 8.f;
				if (is_land && false)
					velocity_(x, y) -= friction_factor * velocity_(x, y);// * math::length(velocity_(x, y));
			}
		}

		// Solve Poisson equation for pressure
		for (int iteration = 0; iteration < 16; ++iteration)
		{
			for (int y = 1; y < N - 1; ++y)
			{
				for (int x = 0; x < N; ++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)
		{
			// 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);
		}
		// 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 = 0; x < N; ++x)
			{
				// Pressure gradient
				math::vector gradient{pressure_(wrap(x + 1), y) - pressure_(wrap(x - 1), y), pressure_(x, y + 1) - pressure_(x, y - 1)};
				velocity_(x, y) -= gradient;
			}
		}

		// Apply boundary conditions for velocity
		for (int i = 0; i < N; ++i)
		{
			// 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];
		}
	}

	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};
		}

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

		for (int y = 0; y < N; ++y)
		{
			for (int x = 0; x < N; ++x)
			{
				if (show_temperature_ || show_pressure_)
				{
					float value = 0.f;
					if (show_temperature_)
						value = 2.f * temperature_(x, y);
					else if (show_pressure_)
						value = 1000.f * pressure_(x, y);

					gfx::color_4f color;
					if (value > 0.f)
					{
						value = 1.f - std::exp(- value);
						color = {value, value * 0.5f, value * 0.125f, 1.f};
					}
					else
					{
						value = 1.f - std::exp(value);
						color = {value * 0.125f, value * 0.5f, value, 1.f};
					}

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

				if (show_velocity_)
				{
					math::point center{x + 0.5f, y + 0.5f};
					auto v = velocity_(x, y);
					if (auto l = math::length(v); l > 0.f)
					{
						float const magnification = 40.f;
						float const max_length = 1.5f;
						v *= 0.5f * max_length * (1.f - std::exp(- magnification * l)) / l;
					}
					auto n = math::ort(v) * 0.3f;

					painter_.triangle(center - v - n, center - v + n, center + v, {255, 255, 255, 255});
				}
			}
		}

		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;
		};

		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 = {} {}", velocity_(x, y)[0], velocity_(x, y)[1]));
			push_text(std::format("P = {}", pressure_(x, y)));
			push_text(std::format("T = {}", temperature_(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_ = true;
	bool show_pressure_ = false;

	util::ndarray<float, 2> obstacle_;
	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> temperature_;
	util::ndarray<float, 2> new_temperature_;
};

namespace psemek::app
{

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

}