#pragma once

#include <psemek/math/matrix.hpp>

namespace psemek::math
{

	template <typename T, std::size_t N>
	struct qr_result
	{
		matrix<T, N, N> q; // orthogonal
		matrix<T, N, N> r; // upper-triangular (for decomposition) or block upper-triangular (for eigenvalues)
	};

	// returns {q, r} such that m = q * r
	template <typename T, std::size_t N>
	qr_result<T, N> qr_decomposition(matrix<T, N, N> const & m)
	{
		matrix<T, N, N> q = matrix<T, N, N>::identity();
		matrix<T, N, N> r = m;

		// Iterate columns
		for (std::size_t i = 0; i < N; ++i)
		{
			// Iterate rows
			for (std::size_t j = i + 1; j < N; ++j)
			{
				T const max = std::max(std::abs(r[i][i]), std::abs(r[j][i]));

				if (max == T{})
					continue;

				// Compute Givens rotation
				T const l = max * std::sqrt(sqr(r[i][i] / max) + sqr(r[j][i] / max));
				T const c =   r[i][i] / l;
				T const s = - r[j][i] / l;

				// Apply rotation to rows of r
				for (std::size_t k = i; k < N; ++k)
				{
					T const x = r[i][k];
					T const y = r[j][k];

					r[i][k] = x * c - y * s;
					r[j][k] = x * s + y * c;
				}

				// Apply inverse rotation to columns of q
				for (std::size_t k = 0; k < N; ++k)
				{
					T const x = q[k][i];
					T const y = q[k][j];

					q[k][i] = x * c - y * s;
					q[k][j] = x * s + y * c;
				}
			}
		}

		return {q, r};
	}

	// returns {q, r} such that m = q * r * q^(-1)
	// NB: q^(-1) = q^T
	template <typename T, std::size_t N>
	qr_result<T, N> qr_eigenvalues(matrix<T, N, N> const & m, std::size_t const iterations)
	{
		matrix<T, N, N> r = m;
		matrix<T, N, N> q = matrix<T, N, N>::identity();

		for (std::size_t iteration = 0; iteration < iterations; ++iteration)
		{
			auto qr = qr_decomposition(r);
			r = qr.r * qr.q;
			q = q * qr.q;
		}

		return {q, r};
	}

}
