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

using namespace psemek;

struct water_1d_app
	: app::app
{
	water_1d_app();

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

	void on_resize(int width, int height) override;

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

	std::size_t const N = 50;
	std::vector<float> bed;
	std::vector<float> height;
	std::vector<float> discharge;

	float x(std::size_t i) const;

	gfx::painter painter;
};

water_1d_app::water_1d_app()
	: app("Water 1D simulation", 4)
{
	simulation_area[0] = {-1.f, 1.f};
	simulation_area[1] = {0.f, 1.f};

	bed.assign(N, 0.f);
	height.assign(N, 0.f);
	discharge.assign(N + 1, 0.f);

	for (std::size_t i = 0; i < N; ++i)
	{
		bed[i] += (0.5f * (N - i)) / N;
		bed[i] += (100.f / (N - i)) / N;
//		bed[i] += (0.25f * i) / N;
//		height[i] += (0.3f * i) / N;
//		height[i] = 0.3f - bed[i];
//		height[i] = 0.25f;
	}

//	bed[N * 3 / 4] = 0.2f;

	discharge[0] = 0.04f;
	discharge[N] = 0.f;
}

void water_1d_app::update()
{
	float const dt = 0.01f;
	float const dx = simulation_area[0].length() / N;
	float const g = 10.f;
	float const max_speed = 1.f;
	log::info() << max_speed << " " << (dx / dt);

	std::vector<float> dh(N, 0.f);
	std::vector<float> dq(N + 1, 0.f);

	for (std::size_t i = 0; i < N; ++i)
	{
		dh[i] = - (discharge[i + 1] - discharge[i]) / dx;
	}

	for (std::size_t i = 0; i < N; ++i)
	{
		height[i] += dh[i] * dt;
	}

	dh.assign(N, 0.f);

	for (std::size_t i = 0; i < N; ++i)
	{
		if (i != 0 && height[i] + bed[i] > bed[i - 1] && height[i - 1] + bed[i - 1] > bed[i]) dh[i] += (height[i - 1] + bed[i - 1] - height[i] - bed[i]);
		if (i != N - 1 && height[i] + bed[i] > bed[i + 1] && height[i + 1] + bed[i + 1] > bed[i]) dh[i] += (height[i + 1] + bed[i + 1] - height[i] - bed[i]);
	}

	for (std::size_t i = 0; i < N; ++i)
	{
		height[i] += dh[i] * 0.1f;
		height[i] = std::max(height[i], 0.f);
	}

	auto hu2 = [this](std::size_t i)
	{
		if (i == 0) return 0.f;
		if (i == N) return 0.f;

		if (height[i] == 0.f) return 0.f;

		float q = (discharge[i] + discharge[i + 1]) / 2.f;
		return q * q / height[i];
	};

	auto gh2 = [this, g](std::size_t i)
	{
		auto v = height[i] + bed[i];
		return 0.5f * g * v * v;
	};

	for (std::size_t i = 1; i < N; ++i)
	{
		dq[i] = -(hu2(i) + gh2(i) - hu2(i - 1) - gh2(i - 1)) / dx;
	}

	for (std::size_t i = 1; i < N; ++i)
	{
		discharge[i] += dq[i] * dt;
//		discharge[i] = (discharge[i - 1] + discharge[i + 1]) / 2.f + dq[i] * dt;
//		discharge[i] -= discharge[i] * dt;
		if (discharge[i] > 0.f)
			discharge[i] = std::min(discharge[i], height[i - 1] * max_speed);
		else
			discharge[i] = -std::min(-discharge[i], height[i] * max_speed);

//		if (height[i - 1] + bed[i - 1] < bed[i])
//			discharge[i] = std::min(discharge[i], 0.f);

//		if (height[i] + bed[i] < bed[i - 1])
//			discharge[i] = std::max(0.f, discharge[i]);
	}
}

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

	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();

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

	for (std::size_t i = 0; i < N; ++i)
	{
		painter.rect({{{x(i), x(i + 1)}, {simulation_area[1].min, simulation_area[1].min + bed[i]}}}, gfx::dark(gfx::yellow).as_color_rgba());
		painter.rect({{{x(i), x(i + 1)}, {simulation_area[1].min + bed[i], simulation_area[1].min + bed[i] + height[i]}}}, gfx::blue);
	}

	for (std::size_t i = 1; i < N; ++i)
	{
		painter.line({x(i), 0.f}, {x(i), discharge[i]}, 0.005f, gfx::red);
	}

	painter.render(transform);
}

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

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

float water_1d_app::x(std::size_t i) const
{
	return geom::lerp(simulation_area[0], (1.f * i) / N);
}

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