#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/log/log.hpp>

using namespace psemek;

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

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

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

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

	std::size_t const N = 100;
	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(options const &, context const &)
{
	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.125f * std::exp(- 1000.f * math::sqr(i * 1.f / N - 0.5f));

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

//	height[0] = 20.f;
}

void water_1d_app::update()
{
	float const dt = 0.005f;
	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);

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

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

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

	gfx::color_rgba const earth_color{127, 63, 0, 255};
	gfx::color_rgba const water_color{0, 0, 255, 255};

	for (std::size_t i = 0; i + 1 < N; ++i)
	{
		math::point p0{x(i + 0) * 0.5f + x(i + 1) * 0.5f, simulation_area[1].min};
		math::point p1{x(i + 1) * 0.5f + x(i + 2) * 0.5f, simulation_area[1].min};

		math::point q0{p0[0], simulation_area[1].min + bed[i]};
		math::point q1{p1[0], simulation_area[1].min + bed[i + 1]};

		math::point w0{p0[0], simulation_area[1].min + bed[i] + height[i]};
		math::point w1{p1[0], simulation_area[1].min + bed[i + 1] + height[i + 1]};

		painter.triangle(p0, p1, q0, earth_color);
		painter.triangle(q0, p1, q1, earth_color);

		painter.triangle(q0, q1, w0, water_color);
		painter.triangle(w0, q1, w1, water_color);
	}

//	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_event(app::resize_event const & event)
{
	app::application_base::on_event(event);

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

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

namespace psemek::app
{

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

}
