#pragma once

#include <psemek/math/matrix.hpp>

namespace psemek::math
{

	template <typename T, std::size_t N>
	struct cholesky_ldl_result
	{
		matrix<T, N, N> l; // lower unit-triangular
		vector<T, N> d;
	};

	// returns lower-triangular l such that m = l * l^T
	// m must be symmetric positive-definite
	template <typename T, std::size_t N>
	matrix<T, N, N> cholesky(matrix<T, N, N> const & m)
	{
		matrix<T, N, N> l = matrix<T, N, N>::zero();

		for (std::size_t i = 0; i < N; ++i)
		{
			for (std::size_t j = 0; j <= i; ++j)
			{
				T sum = T{};
				for (std::size_t k = 0; k < j; ++k)
					sum += l[i][k] * l[j][k];

				if (i == j)
					l[i][j] = std::sqrt(std::max(T{}, sum));
				else
					l[i][j] = (m[i][j] - sum) / l[j][j];
			}
		}

		return l;
	}

	// returns lower unit-triangular l and a diagonal d such that m = l * d * l^T
	// m must be symmetric
	template <typename T, std::size_t N>
	cholesky_ldl_result<T, N> cholesky_ldl(matrix<T, N, N> const & m)
	{
		matrix<T, N, N> l = matrix<T, N, N>::identity();
		vector<T, N> d = vector<T, N>::zero();

		for (std::size_t i = 0; i < N; ++i)
		{
			for (std::size_t j = 0; j <= i; ++j)
			{
				T sum = T{};
				for (std::size_t k = 0; k < j; ++k)
					sum += l[i][k] * l[j][k] * d[k];

				if (i == j)
					d[i] = m[i][i] - sum;
				else
					l[i][j] = (m[i][j] - sum) / d[j];
			}
		}

		return {l, d};
	}

}
