#include <psemek/app/app.hpp>
#include <psemek/app/main.hpp>
#include <psemek/gfx/painter.hpp>
#include <psemek/geom/orthographic.hpp>
#include <psemek/geom/gauss.hpp>

using namespace psemek;

static std::pair<int, int> unpack_triangular_id(int id)
{
	int i = int(std::floor(0.5f * (sqrt(1.f + 8.f * id) - 1.f)));
	int j = id - (i * (i + 1)) / 2;
	return {i, j};
}

static int pack_triangular_id(std::pair<int, int> const & p)
{
	return (p.first * (p.first + 1)) / 2 + p.second;
}

struct water_2d_app
	: app::app
{
	water_2d_app();

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

	void on_resize(int width, int height) override;

	geom::box<float, 2> simulation_area;
	float aspect_ratio;

	int const N = 25;
	int const first_type_count = (N * (N + 1)) / 2;

	std::vector<float> water_height;
	std::vector<geom::vector<float, 2>> flux[3];

	geom::point<float, 2> p[3];

	gfx::painter painter;
};

water_2d_app::water_2d_app()
	: app("Water 2D simulation", 4)
{
	p[0] = {- N / 2.f, 0.f};
	p[1] = {  N / 2.f, 0.f};
	p[2] = {0.f, N * std::sqrt(0.75f)};

	simulation_area |= p[0];
	simulation_area |= p[1];
	simulation_area |= p[2];

	water_height.assign(N * N, 0.5f);
	flux[0].assign(first_type_count, geom::vector{0.f, -5.f});
	flux[1].assign(first_type_count, geom::vector{0.f, -5.f});
	flux[2].assign(first_type_count, geom::vector{0.f, -5.f});

	for (int i = 0; i < N; ++i)
	{
		flux[0][first_type_count - i - 1] = geom::vector<float, 2>::zero();
		flux[1][first_type_count - i - 1] = geom::vector<float, 2>::zero();
		flux[2][first_type_count - i - 1] = geom::vector<float, 2>::zero();
	}
}

void water_2d_app::update()
{
	float const dt = 0.01f;
	float const max_speed = 10.f;

	geom::vector<float, 2> const flux_normal[3] =
	{
		geom::vector{ std::sqrt(3.f) / 2.f, 1.f / 2.f },
		geom::vector{ - std::sqrt(3.f) / 2.f, 1.f / 2.f },
		geom::vector{ 0.f, -1.f },
	};

	auto const rot1 = [this](std::pair<int, int> const & p)
	{
		return std::pair{N - p.second, p.first - p.second};
	};

	auto const rot2 = [rot1](std::pair<int, int> const & p)
	{
		return rot1(rot1(p));
	};

	std::vector<float> dh(water_height.size(), 0.f);
	std::vector<geom::vector<float, 2>> du[3];
	for (int s = 0; s < 3; ++s)
		du[s].assign(first_type_count, geom::vector{0.f, 0.f});

	for (int id = 0; id < N * N; ++id)
	{
		bool const first_type = (id < first_type_count);
		auto const [i, j] = first_type ? unpack_triangular_id(id) : unpack_triangular_id(id - first_type_count);

		int e[3];

		if (first_type)
		{
			e[0] = pack_triangular_id({i, j});
			e[1] = pack_triangular_id(rot1({i + 1, j + 1}));
			e[2] = pack_triangular_id(rot2({i + 1, j}));
		}
		else
		{
			e[0] = pack_triangular_id({i, j});
			e[1] = pack_triangular_id(rot1({i + 2, j + 2}));
			e[2] = pack_triangular_id(rot2({i + 2, j}));
		}

		float const flux_sign = first_type ? -1.f : 1.f;

		for (int s = 0; s < 3; ++s)
			dh[id] += flux_sign * geom::dot(flux_normal[s], flux[s][e[s]]) * dt;
	}

	for (int id = 0; id < N * N; ++id)
	{
		water_height[id] += dh[id];
		water_height[id] = std::max(0.f, water_height[id]);
	}

	for (int s = 0; s < 3; ++s)
	{
		for (int id = 0; id < first_type_count - N; ++id)
		{
			auto const [i, j] = unpack_triangular_id(id);

			std::pair<int, int> tin, tout;

			if (s == 0)
			{
				tout = {i, j};
				tin = {i, j};
			}
			else if (s == 1)
			{
				tout = rot2({i + 1, j});
				tin = rot2({i + 2, j});
			}
			else // if (s == 2)
			{
				tout = rot1({i + 1, j + 1});
				tin = rot1({i + 2, j + 2});
			}


		}
	}

	for (int s = 0; s < 3; ++s)
	{
		for (int id = 0; id < first_type_count - N; ++id)
		{
			flux[s][id] += du[s][id];

			auto const [i, j] = unpack_triangular_id(id);

			std::pair<int, int> tid;
			bool first_type;

			float const nflux = geom::dot(flux[s][id], flux_normal[s]);
			bool const forward = nflux > 0.f;

			if (forward)
			{
				first_type = true;
				if (s == 0)
					tid = {i, j};
				else if (s == 1)
					tid = rot2({i + 1, j});
				else // if (s == 2)
					tid = rot1({i + 1, j + 1});
			}
			else
			{
				first_type = false;
				if (s == 0)
					tid = {i, j};
				else if (s == 1)
					tid = rot2({i + 2, j});
				else // if (s == 2)
					tid = rot1({i + 2, j + 2});
			}

			auto const h = water_height[pack_triangular_id(tid) + (first_type ? 0 : first_type_count)];
			float const max_flux = h * max_speed;

			if (forward && nflux > max_flux)
			{
				flux[s][id] -= flux_normal[s] * (nflux - max_flux);
			}

			if (!forward && (-nflux > max_flux))
			{
				flux[s][id] += flux_normal[s] * ((-nflux) - max_flux);
			}

//			if (forward)
//				flux[s][id] = std::min(flux[s][id], h * max_speed);
//			else
//				flux[s][id] = std::max(flux[s][id], - h * max_speed);
		}
	}
}

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

	gl::Enable(gl::BLEND);
	gl::BlendFunc(gl::SRC_ALPHA, gl::ONE_MINUS_SRC_ALPHA);

	geom::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.1f;
		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;
	}

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

	float const pixel = 2.f / std::sqrt(width() * height() * geom::det(transform));

	// Background
	painter.triangle(p[0], p[1], p[2], gfx::white);

	// Water tiles
	for (int id = 0; id < N * N; ++id)
	{
		bool const first_type = (id < first_type_count);
		auto const [i, j] = first_type ? unpack_triangular_id(id) : unpack_triangular_id(id - first_type_count);
		(void)pack_triangular_id;

		auto at = [this](int i, int j)
		{
			float const u = (i * 1.f) / N;
			float const v = (j * 1.f) / N;
			return geom::lerpn(p[0], 1.f - u, p[1], v, p[2], u - v);
		};

		auto f = [](float h)
		{
			float const C = 1.f;
			float const e = std::exp(- C * h);
			return (1.f - e) / (1.f + e);
		};

		auto color = gfx::to_coloru8(gfx::color_4f{0.f, 0.f, 1.f, f(water_height[id])});

		if (water_height[id] < 0.f)
			color = gfx::magenta;

		if (first_type)
		{
			painter.triangle(at(i, j), at(i + 1, j + 1), at(i + 1, j), color);
		}
		else
		{
			painter.triangle(at(i + 1, j), at(i + 1, j + 1), at(i + 2, j + 1), color);
		}
	}

	// Edges
//	if(false)
	for (int i = 1; i <= N; ++i)
	{
		float t = (i * 1.f) / N;
		painter.line(geom::lerp(p[0], p[1], t), geom::lerp(p[0], p[2], t), 2.f * pixel, gfx::black, false);
		painter.line(geom::lerp(p[1], p[0], t), geom::lerp(p[1], p[2], t), 2.f * pixel, gfx::black, false);
		painter.line(geom::lerp(p[2], p[0], t), geom::lerp(p[2], p[1], t), 2.f * pixel, gfx::black, false);
	}

	// Flux
//	if(false)
	for (int side = 0; side < 3; ++side)
	{
		auto at = [this, side](int i, int j)
		{
			float const u = (i * 1.f) / N;
			float const v = (j * 1.f) / N;
			return geom::lerpn(p[side], 1.f - u, p[(side + 1) % 3], v, p[(side + 2) % 3], u - v);
		};

		for (int id = 0; id < first_type_count; ++id)
		{
			float const scale = 0.1f;

			auto const [i, j] = unpack_triangular_id(id);

			auto const p0 = at(i + 1, j);
			auto const p1 = at(i + 1, j + 1);

			auto const p = geom::lerp(p0, p1, 0.5f);

			painter.line(p, p + flux[side][id] * scale, pixel * 2.f, gfx::red, false);
		}
	}

	painter.render(transform);
}

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

	aspect_ratio = (width * 1.f) / height;
}

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