#pragma once

#include <psemek/util/ndarray.hpp>
#include <psemek/math/math.hpp>
#include <psemek/math/dual.hpp>

namespace psemek::phys
{

	template <typename T>
	std::pair<std::size_t, T> shallow_water_steady_state_single_step(util::ndarray<T, 2> const & depth, util::ndarray<T, 2> & water_height,
		util::ndarray<T, 2> & x_velocity, util::ndarray<T, 2> & y_velocity,
		T gravity, std::size_t max_iterations, T max_error)
	{
		std::size_t const n = depth.width() - 2;
		std::size_t const m = depth.height() - 2;

		util::ndarray<T, 2> delta_height({n + 2, m + 2}, 0.f);
		util::ndarray<T, 2> delta_x_velocity({n + 1, m + 2}, 0.f);
		util::ndarray<T, 2> delta_y_velocity({n + 2, m + 1}, 0.f);

		auto h_dual = [&](auto i, auto j) { return math::dual<float, 2>{water_height(i, j), {delta_height(i, j), 0.f}}; };
		auto ux_dual = [&](auto i, auto j) { return math::dual<float, 2>{x_velocity(i, j + 1), {delta_x_velocity(i, j + 1), 0.f}}; };
		auto uy_dual = [&](auto i, auto j) { return math::dual<float, 2>{y_velocity(i + 1, j), {delta_y_velocity(i + 1, j), 0.f}}; };

		auto h_diag = [&](auto i, auto j) { return math::dual<float, 2>{water_height(i, j), {0.f, 1.f}}; };
		auto ux_diag = [&](auto i, auto j) { return math::dual<float, 2>{x_velocity(i, j + 1), {0.f, 1.f}}; };
		auto uy_diag = [&](auto i, auto j) { return math::dual<float, 2>{y_velocity(i + 1, j), {0.f, 1.f}}; };

		T error = 0;

		T omega = T{0.1};

		std::size_t iteration = 0;
		for (; iteration < max_iterations; ++iteration)
		{
			// Single iteration of Gauss-Seidel
			error = 0;

			// Mass conservation equations
			for (std::size_t j = 1; j < m + 1; ++j)
			{
				for (std::size_t i = 1; i < n + 1; ++i)
				{
					math::dual<float, 2> value =
						+ (h_dual(i + 1, j    ) + h_diag(i, j)) * ux_dual(i    , j - 1) / 2.f
						- (h_dual(i - 1, j    ) + h_diag(i, j)) * ux_dual(i - 1, j - 1) / 2.f
						+ (h_dual(i    , j + 1) + h_diag(i, j)) * uy_dual(i - 1, j    ) / 2.f
						- (h_dual(i    , j - 1) + h_diag(i, j)) * uy_dual(i - 1, j - 1) / 2.f
						;

//					error += math::sqr(value.scalar - value.delta[0] - value.delta[1] * delta_height(i, j));
					error += math::sqr(value.scalar);

					if (value.delta[1] != 0.f)
						delta_height(i, j) = math::lerp(delta_height(i, j), (-value.scalar - value.delta[0]) / value.delta[1], omega);
				}
			}

			// X-momentum conservation equations
			for (std::size_t j = 0; j < m; ++j)
			{
				for (std::size_t i = 1; i < n; ++i)
				{
					math::dual<float, 2> value =
						+ h_dual(i + 1, j + 1) * (math::sqr(ux_dual(i + 1, j    ) + math::sqr(ux_diag(i, j)))) / 2.f
						- h_dual(i    , j + 1) * (math::sqr(ux_dual(i - 1, j    ) + math::sqr(ux_diag(i, j)))) / 2.f

						+ (h_dual(i, j + 1) + h_dual(i + 1, j + 1) + h_dual(i, j + 2) + h_dual(i + 1, j + 2))
							* (ux_diag(i, j) + ux_dual(i, j + 1)) * (uy_dual(i - 1, j + 1) + uy_dual(i, j + 1)) / 16.f

						- (h_dual(i, j + 1) + h_dual(i + 1, j + 1) + h_dual(i, j) + h_dual(i - 1, j))
							* (ux_diag(i, j) + ux_dual(i, j - 1)) * (uy_dual(i - 1, j) + uy_dual(i, j)) / 16.f

						+ gravity * (math::sqr(h_dual(i + 1, j + 1)) - math::sqr(h_dual(i, j + 1))) / 2.f

						- gravity * (h_dual(i + 1, j + 1) + h_dual(i, j + 1)) * (depth(i + 1, j + 1) - depth(i, j + 1)) / 2.f
						;

					value = 0;

//					error += math::sqr(value.scalar - value.delta[0] - value.delta[1] * delta_x_velocity(i, j + 1));
					error += math::sqr(value.scalar);

					if (value.delta[1] != 0.f)
						delta_x_velocity(i, j + 1) = math::lerp(delta_x_velocity(i, j + 1), (value.scalar - value.delta[0]) / value.delta[1], omega);
				}
			}

			// Y-momentum conservation equations
			for (std::size_t j = 1; j < m; ++j)
			{
				for (std::size_t i = 0; i < n; ++i)
				{
					math::dual<float, 2> value =
						+ h_dual(i + 1, j + 1) * (math::sqr(uy_dual(i    , j + 1) + math::sqr(uy_diag(i, j)))) / 2.f
						- h_dual(i + 1, j    ) * (math::sqr(uy_dual(i    , j - 1) + math::sqr(uy_diag(i, j)))) / 2.f

						+ (h_dual(i + 1, j) + h_dual(i + 1, j + 1) + h_dual(i + 2, j) + h_dual(i + 2, j + 1))
							* (uy_diag(i, j) + uy_dual(i + 1, j)) * (ux_dual(i + 1, j - 1) + ux_dual(i + 1, j)) / 16.f

						- (h_dual(i + 1, j) + h_dual(i + 1, j + 1) + h_dual(i, j) + h_dual(i, j - 1))
							* (uy_diag(i, j) + uy_dual(i - 1, j)) * (ux_dual(i, j - 1) + ux_dual(i, j)) / 16.f

						+ gravity * (math::sqr(h_dual(i + 1, j + 1)) - math::sqr(h_dual(i + 1, j))) / 2.f

						- gravity * (h_dual(i + 1, j + 1) + h_dual(i + 1, j)) * (depth(i + 1, j + 1) - depth(i + 1, j)) / 2.f
						;

					value = 0;

//					error += math::sqr(value.scalar - value.delta[0] - value.delta[1] * delta_y_velocity(i + 1, j));
					error += math::sqr(value.scalar);

					if (value.delta[1] != 0.f)
						delta_y_velocity(i + 1, j) = math::lerp(delta_y_velocity(i + 1, j), (-value.scalar - value.delta[0]) / value.delta[1], omega);
				}
			}

			int temp = 42;
			(void)temp;

			(void)max_error;
		}

		for (std::size_t j = 0; j < m + 2; ++j)
			for (std::size_t i = 0; i < n + 2; ++i)
				water_height(i, j) += delta_height(i, j);

		for (std::size_t j = 0; j < m + 2; ++j)
			for (std::size_t i = 0; i < n + 1; ++i)
				x_velocity(i, j) += delta_x_velocity(i, j);

		for (std::size_t j = 0; j < m + 1; ++j)
			for (std::size_t i = 0; i < n + 2; ++i)
				y_velocity(i, j) += delta_y_velocity(i, j);

		return {iteration, error};
	}

}
