#include <psemek/test/test.hpp>

#include <psemek/geom/matrix.hpp>
#include <psemek/geom/gauss.hpp>
#include <psemek/geom/qr.hpp>
#include <psemek/geom/cholesky.hpp>
#include <psemek/random/generator.hpp>
#include <psemek/random/normal.hpp>

using namespace psemek;
using namespace psemek::geom;

test_case(geom_matrix_empty)
{
	matrix<int, 0, 0> m00;
	matrix<int, 1, 0> m10;
	matrix<int, 0, 1> m01;

	static_assert(std::is_same_v<decltype(m00)::scalar_type, int>);
	static_assert(std::is_same_v<decltype(m10)::scalar_type, int>);
	static_assert(std::is_same_v<decltype(m01)::scalar_type, int>);

	expect_equal(m00.columns(), 0);
	expect_equal(m00.rows(), 0);
	expect_equal(m10.columns(), 0);
	expect_equal(m10.rows(), 1);
	expect_equal(m01.columns(), 1);
	expect_equal(m01.rows(), 0);

	expect_throw(m00[0][0], detail::empty_array_exception);
	expect_throw(m10[0][0], detail::empty_array_exception);
	expect_throw(m01[0][0], detail::empty_array_exception);
}

test_case(geom_matrix_init)
{
	matrix<int, 2, 2> m;
	m[0][0] = 1;
	m[0][1] = 2;
	m[1][0] = 3;
	m[1][1] = 4;

	expect_equal(m.columns(), 2);
	expect_equal(m.rows(), 2);
	expect_equal(m[0][0], 1);
	expect_equal(m[0][1], 2);
	expect_equal(m[1][0], 3);
	expect_equal(m[1][1], 4);
}

test_case(geom_matrix_by__rows)
{
	matrix<int, 2, 2> m;
	m[0][0] = 1;
	m[0][1] = 2;
	m[1][0] = 3;
	m[1][1] = 4;

	expect_equal(m, (by_rows(vector{1, 2}, vector{3, 4})));
}

test_case(geom_matrix_by__columns)
{
	matrix<int, 2, 2> m;
	m[0][0] = 1;
	m[0][1] = 2;
	m[1][0] = 3;
	m[1][1] = 4;

	expect_equal(m, (by_columns(vector{1, 3}, vector{2, 4})));
}

test_case(geom_matrix_zero)
{
	auto m = matrix<float, 5, 5>::zero();

	for (std::size_t i = 0; i < m.rows(); ++i)
	{
		for (std::size_t j = 0; j < m.columns(); ++j)
		{
			expect_equal(m[i][j], 0.f);
		}
	}
}

test_case(geom_matrix_identity)
{
	auto m = matrix<float, 5, 5>::identity();

	for (std::size_t i = 0; i < m.rows(); ++i)
	{
		for (std::size_t j = 0; j < m.columns(); ++j)
		{
			expect_equal(m[i][j], (i == j ? 1.f : 0.f));
		}
	}
}

test_case(geom_matrix_scalar)
{
	auto m = matrix<float, 5, 5>::scalar(3.f);

	for (std::size_t i = 0; i < m.rows(); ++i)
	{
		for (std::size_t j = 0; j < m.columns(); ++j)
		{
			expect_equal(m[i][j], (i == j ? 3.f : 0.f));
		}
	}
}

test_case(geom_matrix_arithmetic)
{
	vector v{1.f, 2.f};

	matrix<float, 2, 2> m1, m2;

	m1[0][0] = 1.f;
	m1[0][1] = 2.f;
	m1[1][0] = 3.f;
	m1[1][1] = 4.f;

	m2[0][0] = 5.f;
	m2[0][1] = 6.f;
	m2[1][0] = 7.f;
	m2[1][1] = 8.f;

	expect_equal(m1 * v, (vector{5.f, 11.f}));
	expect_equal(m2 * v, (vector{17.f, 23.f}));

	expect_equal(-m1, (by_rows(vector{-1.f, -2.f}, vector{-3.f, -4.f})));
	expect_equal(5.f * m1, (by_rows(vector{5.f, 10.f}, vector{15.f, 20.f})));
	expect_equal(m1 * 5.f, (by_rows(vector{5.f, 10.f}, vector{15.f, 20.f})));
	expect_equal(m1 + m2, (by_rows(vector{6.f, 8.f}, vector{10.f, 12.f})));
	expect_equal(m1 - m2, (by_rows(vector{-4.f, -4.f}, vector{-4.f, -4.f})));
	expect_equal(m1 * m2, (by_rows(vector{19.f, 22.f}, vector{43.f, 50.f})));
	expect_equal(m2 * m1, (by_rows(vector{23.f, 34.f}, vector{31.f, 46.f})));

	expect_equal(transpose(m1), (by_rows(vector{1.f, 3.f}, vector{2.f, 4.f})));
	expect_equal(transpose(transpose(m1)), m1);
}

test_case(geom_matrix_trace)
{
	matrix<float, 2, 2> m;

	m[0][0] = 1.f;
	m[0][1] = 2.f;
	m[1][0] = 3.f;
	m[1][1] = 4.f;

	expect_close(trace(m), 5.f, 1e-6);
}

test_case(geom_matrix_det)
{
	matrix<float, 2, 2> m;

	m[0][0] = 1.f;
	m[0][1] = 2.f;
	m[1][0] = 3.f;
	m[1][1] = 4.f;

	expect_close(det(m), -2.f, 1e-6);
}

test_case(geom_matrix_solve)
{
	matrix<float, 2, 2> m;

	m[0][0] = 1.f;
	m[0][1] = 2.f;
	m[1][0] = 3.f;
	m[1][1] = 4.f;

	vector v{1.f, 2.f};

	auto s = solve(m, v);

	expect(s);
	expect_small(length(*s - vector{0.f, 0.5f}), 1e-6);
}

test_case(geom_matrix_inverse)
{
	matrix<float, 2, 2> m;

	m[0][0] = 1.f;
	m[0][1] = 2.f;
	m[1][0] = 3.f;
	m[1][1] = 4.f;

	auto i = inverse(m);

	expect(i);
	expect_small(frobenius_norm(*i - by_rows(vector{-2.f, 1.f}, vector{1.5f, -0.5f})), 1e-6);
}

test_case(geom_matrix_solve__lower__triangular)
{
	random::generator rng;
	random::normal_distribution<float> normal;

	for (int iteration = 0; iteration < 1024; ++iteration)
	{
		matrix<float, 8, 8> m = matrix<float, 8, 8>::zero();

		for (std::size_t i = 0; i < m.rows(); ++i)
			for (std::size_t j = 0; j <= i; ++j)
				m[i][j] = (i == j) ? (1.f + std::abs(normal(rng))) : normal(rng);

		vector<float, 8> v;

		for (std::size_t i = 0; i < v.dimension(); ++i)
			v[i] = normal(rng);

		auto u = v;
		solve_lower_triangular(m, u);

		expect_small(length(m * u - v), 1e-5);
	}
}

test_case(geom_matrix_solve__upper__triangular)
{
	random::generator rng;
	random::normal_distribution<float> normal;

	for (int iteration = 0; iteration < 1024; ++iteration)
	{
		matrix<float, 8, 8> m = matrix<float, 8, 8>::zero();

		for (std::size_t i = 0; i < m.rows(); ++i)
			for (std::size_t j = i; j < m.columns(); ++j)
				m[i][j] = (i == j) ? (1.f + std::abs(normal(rng))) : normal(rng);

		vector<float, 8> v;

		for (std::size_t i = 0; i < v.dimension(); ++i)
			v[i] = normal(rng);

		auto u = v;
		solve_upper_triangular(m, u);

		expect_small(length(m * u - v), 1e-5);
	}
}

test_case(geom_matrix_inverse__lower__triangular)
{
	random::generator rng;
	random::normal_distribution<float> normal;

	for (int iteration = 0; iteration < 1024; ++iteration)
	{
		matrix<float, 8, 8> m = matrix<float, 8, 8>::zero();

		for (std::size_t i = 0; i < m.rows(); ++i)
			for (std::size_t j = 0; j <= i; ++j)
				m[i][j] = (i == j) ? (1.f + std::abs(normal(rng))) : normal(rng);

		auto inv = inverse_lower_triangular(m);

		expect(inv);
		expect_small(frobenius_norm(m.identity() - m * *inv), 1e-4);
	}
}

test_case(geom_matrix_inverse__upper__triangular)
{
	random::generator rng;
	random::normal_distribution<float> normal;

	for (int iteration = 0; iteration < 1024; ++iteration)
	{
		matrix<float, 8, 8> m = matrix<float, 8, 8>::zero();

		for (std::size_t i = 0; i < m.rows(); ++i)
			for (std::size_t j = i; j < m.columns(); ++j)
				m[i][j] = (i == j) ? (1.f + std::abs(normal(rng))) : normal(rng);

		auto inv = inverse_upper_triangular(m);

		expect(inv);
		expect_small(frobenius_norm(m.identity() - m * *inv), 1e-4);
	}
}

test_case(geom_matrix_qr)
{
	random::generator rng;
	random::normal_distribution<float> normal;

	for (int iteration = 0; iteration < 1024; ++iteration)
	{
		matrix<float, 8, 8> m;
		for (std::size_t i = 0; i < m.rows(); ++i)
			for (std::size_t j = 0; j < m.columns(); ++j)
				m[i][j] = normal(rng);

		auto result = qr_decomposition(m);

		expect_small(frobenius_norm(m - result.q * result.r), 1e-5);
		expect_small(frobenius_norm(matrix<float, 8, 8>::identity() - result.q * transpose(result.q)), 1e-5);

		for (std::size_t i = 0; i < m.rows(); ++i)
			for (std::size_t j = i + 1; j < m.columns(); ++j)
				expect_small(result.r[j][i], 1e-6);
	}
}

test_case(geom_matrix_qr__eig_symmetric)
{
	random::generator rng;
	random::normal_distribution<float> normal;

	for (int iteration = 0; iteration < 64; ++iteration)
	{
		matrix<float, 8, 8> m;
		for (std::size_t i = 0; i < m.rows(); ++i)
			for (std::size_t j = 0; j < m.columns(); ++j)
				m[i][j] = normal(rng);

		m = transpose(m) * m;

		auto result = qr_eigenvalues(m, 1024);

		expect_small(frobenius_norm(m * result.q - result.q * result.r), 1e-3);
		expect_small(frobenius_norm(matrix<float, 8, 8>::identity() - result.q * transpose(result.q)), 1e-4);

		for (std::size_t i = 0; i < m.rows(); ++i)
			for (std::size_t j = i + 1; j < m.columns(); ++j)
				expect_small(result.r[j][i], 1e-3);

		for (std::size_t i = 0; i < m.columns(); ++i)
		{
			auto v = column(result.q, i);
			expect_small(length(m * v - result.r[i][i] * v), 1e-3);
		}
	}
}

test_case(geom_matrix_qr__eig_general)
{
	random::generator rng;
	random::normal_distribution<float> normal;

	for (int iteration = 0; iteration < 64; ++iteration)
	{
		matrix<float, 8, 8> m;
		for (std::size_t i = 0; i < m.rows(); ++i)
			for (std::size_t j = 0; j < m.columns(); ++j)
				m[i][j] = normal(rng);

		auto result = qr_eigenvalues(m, 4096);

		expect_small(frobenius_norm(m * result.q - result.q * result.r), 1e-2);
		expect_small(frobenius_norm(matrix<float, 8, 8>::identity() - result.q * transpose(result.q)), 1e-3);

		for (std::size_t i = 0; i < m.rows(); ++i)
			for (std::size_t j = i + 2; j < m.columns(); ++j)
				expect_small(result.r[j][i], 1e-2);

		for (std::size_t i = 0; i < m.columns(); ++i)
		{
			// Don't check complex eigenvalues
			if (i > 0 && std::abs(result.r[i][i - 1]) > 1e-3f) continue;
			if (i + 1 < m.rows() && std::abs(result.r[i + 1][i]) > 1e-3f) continue;

			// TODO: properly check eigenvalues
//			auto v = column(result.q, i);
//			expect_small(length(m * v - result.r[i][i] * v), 1e-3);
		}
	}
}

test_case(geom_matrix_cholesky)
{
	random::generator rng;
	random::normal_distribution<float> normal;

	for (int iteration = 0; iteration < 64; ++iteration)
	{
		matrix<float, 8, 8> m;
		for (std::size_t i = 0; i < m.rows(); ++i)
			for (std::size_t j = 0; j < m.columns(); ++j)
				m[i][j] = normal(rng);

		m = transpose(m) * m;

		auto l = cholesky(m);

		expect_small(frobenius_norm(m - l * transpose(l)), 1e-6);

		for (std::size_t i = 0; i < m.rows(); ++i)
			for (std::size_t j = i + 1; j < m.columns(); ++j)
				expect_small(l[i][j], 1e-6);
	}
}

test_case(geom_matrix_cholesky__ldl)
{
	random::generator rng;
	random::normal_distribution<float> normal;

	for (int iteration = 0; iteration < 64; ++iteration)
	{
		matrix<float, 8, 8> m;
		for (std::size_t i = 0; i < m.rows(); ++i)
			for (std::size_t j = 0; j < m.columns(); ++j)
				m[i][j] = normal(rng);

		m = (m + transpose(m)) * 0.5f;

		auto ldl = cholesky_ldl(m);

		expect_small(frobenius_norm(m - ldl.l * m.diagonal(ldl.d) * transpose(ldl.l)), 1e-4);

		for (std::size_t i = 0; i < m.rows(); ++i)
			expect_close(ldl.l[i][i], 1.f, 1e-6);

		for (std::size_t i = 0; i < m.rows(); ++i)
			for (std::size_t j = i + 1; j < m.columns(); ++j)
				expect_small(ldl.l[i][j], 1e-6);
	}
}