#pragma once

#include <psemek/math/matrix.hpp>

#include <cmath>

namespace psemek::math
{

	namespace detail
	{

		// A helper to treat vectors & matrices uniformly

		template <typename X>
		struct gauss_helper;

		template <typename X>
		gauss_helper(X &) -> gauss_helper<X>;

		template <typename T, std::size_t N>
		struct gauss_helper<vector<T, N>>
		{
			vector<T, N> & v;

			static constexpr std::size_t columns() { return 1; }
			T & at(std::size_t row, std::size_t) { return v[row]; }
		};

		template <typename T, std::size_t N, std::size_t M>
		struct gauss_helper<matrix<T, N, M>>
		{
			matrix<T, N, M> & m;

			static constexpr std::size_t columns() { return M; }
			T & at(std::size_t row, std::size_t col) { return m[row][col]; }
		};

		template <typename F, typename ... Args>
		void for_each(F && f, Args & ... args)
		{
			(f(args), ...);
		}

	}

	template <typename T, std::size_t N>
	T det(matrix<T, N, N> m)
	{
		using std::abs;

		std::size_t sign = 0;

		for (std::size_t i = 0; i < N; ++i)
		{
			// find maximal modulus along i-th column
			auto M = abs(m[i][i]);
			std::size_t k = i;
			for (std::size_t j = i + 1; j < N; ++j)
			{
				auto const n = abs(m[j][i]);
				if (n > M)
				{
					M = n;
					k = j;
				}
			}

			// avoid dividing by zero
			if (M == T{})
				return T{};

			// swap rows i & k
			if (i != k)
			{
				sign += 1;

				for (std::size_t j = i; j < N; ++j)
					std::swap(m[i][j], m[k][j]);
			}

			// zero out the column under i,i
			for (std::size_t k = i + 1; k < N; ++k)
			{
				auto r = m[k][i] / m[i][i];
				for (std::size_t j = i; j < N; ++j)
				{
					m[k][j] -= m[i][j] * r;
				}
			}
		}

		// determinant is product of diagonal entries + sign
		T res = ((sign % 2) == 0) ? T{1} : -T{1};
		for (std::size_t i = 0; i < N; ++i)
			res *= m[i][i];
		return res;
	}

	// each element of the rhs can be a vector<N> or a matrix<N, K>
	// the set of equations (m * x = rhs)... is solved simultaneously column-wise
	// returns false is the matrix is degenerate
	template <typename T, std::size_t N, typename ... RHS>
	bool gauss(matrix<T, N, N> m, RHS & ... rhs)
	{
		using std::abs;

		// forward elimination
		for (std::size_t i = 0; i < N; ++i)
		{
			// find maximal modulus along i-th column
			auto M = abs(m[i][i]);
			std::size_t k = i;
			for (std::size_t j = i + 1; j < N; ++j)
			{
				auto const n = abs(m[j][i]);
				if (n > M)
				{
					M = n;
					k = j;
				}
			}

			// avoid dividing by zero
			if (M == T{})
				return false;

			// swap rows i & k
			if (i != k)
			{
				for (std::size_t j = i; j < N; ++j)
					std::swap(m[i][j], m[k][j]);

				detail::for_each([i, k](auto & rhs){
					detail::gauss_helper h{rhs};
					for (std::size_t c = 0; c < h.columns(); ++c)
					{
						std::swap(h.at(i, c), h.at(k, c));
					}
				}, rhs...);
			}

			// make i,i equal 1
			{
				auto r = m[i][i];
				for (std::size_t j = i; j < N; ++j)
				{
					m[i][j] /= r;
				}

				detail::for_each([i, r](auto & rhs){
					detail::gauss_helper h{rhs};
					for (std::size_t c = 0; c < h.columns(); ++c)
					{
						h.at(i, c) /= r;
					}
				}, rhs...);
			}

			// zero out the column under i,i
			for (std::size_t k = i + 1; k < N; ++k)
			{
				auto r = m[k][i];
				for (std::size_t j = i; j < N; ++j)
				{
					m[k][j] -= m[i][j] * r;
				}

				detail::for_each([i, k, r](auto & rhs){
					detail::gauss_helper h{rhs};
					for (std::size_t c = 0; c < h.columns(); ++c)
					{
						h.at(k, c) -= h.at(i, c) * r;
					}
				}, rhs...);
			}
		}

		// backward elimination
		for (std::size_t i = N; i --> 0;)
		{
			for (std::size_t j = 0; j < i; ++j)
			{
				auto r = m[j][i];
				m[j][i] = T{};

				detail::for_each([i, j, r](auto & rhs){
					detail::gauss_helper h{rhs};
					for (std::size_t c = 0; c < h.columns(); ++c)
					{
						h.at(j, c) -= r * h.at(i, c);
					}
				}, rhs...);
			}
		}

		return true;
	}

	template <typename T, std::size_t N>
	std::optional<vector<T, N>> solve(matrix<T, N, N> const & m, vector<T, N> v)
	{
		if (!gauss(m, v))
			return std::nullopt;
		return v;
	}

	template <typename T, std::size_t N>
	std::optional<matrix<T, N, N>> inverse(matrix<T, N, N> m)
	{
		matrix<T, N, N> r = matrix<T, N, N>::identity();
		if (!gauss(m, r))
			return std::nullopt;
		return r;
	}

	// Least-squares solution of m*x=b with full-rank m
	template <typename T, std::size_t N, std::size_t M>
	std::optional<vector<T, M>> least_squares(matrix<T, N, M> const & m, vector<T, N> const & b)
	{
		if constexpr (N == M)
		{
			auto rhs = b;
			if (gauss(m, rhs))
				return rhs;
			return {};
		}
		else if constexpr (N < M)
		{
			auto const mt = transpose(m);
			auto rhs = b;
			if (gauss(m * mt, rhs))
				return mt * rhs;
			return {};
		}
		else // if constexpr (N > M)
		{
			auto const mt = transpose(m);
			auto rhs = mt * b;
			if (gauss(mt * m, rhs))
				return rhs;
			return {};
		}
	}

	template <typename T, std::size_t N, typename ... RHS>
	bool solve_lower_triangular(matrix<T, N, N> const & m, RHS & ... rhs)
	{
		for (std::size_t i = 0; i < N; ++i)
		{
			if (m[i][i] == T{}) return false;

			for (std::size_t j = 0; j < i; ++j)
			{
				detail::for_each([&](auto & rhs){
					detail::gauss_helper h{rhs};
					for (std::size_t c = 0; c < h.columns(); ++c)
					{
						h.at(i, c) -= h.at(j, c) * m[i][j];
					}
				}, rhs...);
			}

			detail::for_each([&](auto & rhs){
				detail::gauss_helper h{rhs};
				for (std::size_t c = 0; c < h.columns(); ++c)
				{
					h.at(i, c) /= m[i][i];
				}
			}, rhs...);
		}

		return true;
	}

	template <typename T, std::size_t N, typename ... RHS>
	bool solve_upper_triangular(matrix<T, N, N> const & m, RHS & ... rhs)
	{
		for (std::size_t i = N; i --> 0;)
		{
			if (m[i][i] == T{}) return false;

			for (std::size_t j = i + 1; j < N; ++j)
			{
				detail::for_each([&](auto & rhs){
					detail::gauss_helper h{rhs};
					for (std::size_t c = 0; c < h.columns(); ++c)
					{
						h.at(i, c) -= h.at(j, c) * m[i][j];
					}
				}, rhs...);
			}

			detail::for_each([&](auto & rhs){
				detail::gauss_helper h{rhs};
				for (std::size_t c = 0; c < h.columns(); ++c)
				{
					h.at(i, c) /= m[i][i];
				}
			}, rhs...);
		}

		return true;
	}

	template <typename T, std::size_t N>
	std::optional<matrix<T, N, N>> inverse_lower_triangular(matrix<T, N, N> const & m)
	{
		auto result = matrix<T, N, N>::identity();
		if (!solve_lower_triangular(m, result))
			return std::nullopt;
		return result;
	}

	template <typename T, std::size_t N>
	std::optional<matrix<T, N, N>> inverse_upper_triangular(matrix<T, N, N> const & m)
	{
		auto result = matrix<T, N, N>::identity();
		if (!solve_upper_triangular(m, result))
			return std::nullopt;
		return result;
	}

}
