#pragma once

#include <psemek/math/vector.hpp>
#include <psemek/math/point.hpp>
#include <psemek/math/matrix.hpp>
#include <psemek/math/gauss.hpp>

namespace psemek::math
{

	// Affine transformation from M-dimensional to N-dimensional affine space
	template <typename T, std::size_t N, std::size_t M>
	struct affine_transform
	{
		matrix<T, N, M + 1> m;

		static affine_transform zero();
		static affine_transform identity();

		affine_transform();
		affine_transform(matrix<T, N, M + 1> const & matrix);
		affine_transform(matrix<T, N, M> const & linear, vector<T, N> const & translation);
		affine_transform(affine_transform const &) = default;

		affine_transform & operator = (affine_transform const &) = default;

		matrix<T, N, M + 1> affine_matrix() const;
		matrix<T, N, M> linear_matrix() const;
		vector<T, N> translation_vector() const;
		matrix<T, N + 1, M + 1> homogeneous_matrix() const;

		vector<T, N> operator()(vector<T, M> const & v) const;
		point<T, N> operator()(point<T, M> const & p) const;
	};

	template <typename T, std::size_t N, std::size_t M>
	affine_transform<T, N, M> affine_transform<T, N, M>::zero()
	{
		return affine_transform<T, N, M>{matrix<T, N, M + 1>::zero()};
	}

	template <typename T, std::size_t N, std::size_t M>
	affine_transform<T, N, M> affine_transform<T, N, M>::identity()
	{
		return affine_transform<T, N, M>{matrix<T, N, M + 1>::identity()};
	}

	template <typename T, std::size_t N, std::size_t M>
	affine_transform<T, N, M>::affine_transform()
		: m{matrix<T, N, M+1>::zero()}
	{
		for (std::size_t i = 0; i < std::min(N, M); ++i)
			m[i][i] = T{1};
	}

	template <typename T, std::size_t N, std::size_t M>
	affine_transform<T, N, M>::affine_transform(matrix<T, N, M + 1> const & matrix)
		: m{matrix}
	{}

	template <typename T, std::size_t N, std::size_t M>
	affine_transform<T, N, M>::affine_transform(matrix<T, N, M> const & linear, vector<T, N> const & translation)
	{
		for (std::size_t r = 0; r < N; ++r)
		{
			for (std::size_t c = 0; c < M; ++c)
				m[r][c] = linear[r][c];
			m[r][M] = translation[r];
		}
	}

	template <typename T, std::size_t N, std::size_t M>
	matrix<T, N, M + 1> affine_transform<T, N, M>::affine_matrix() const
	{
		return m;
	}

	template <typename T, std::size_t N, std::size_t M>
	matrix<T, N, M> affine_transform<T, N, M>::linear_matrix() const
	{
		matrix<T, N, M> result;
		for (std::size_t r = 0; r < N; ++r)
		{
			for (std::size_t c = 0; c < M; ++c)
				result[r][c] = m[r][c];
		}
		return result;
	}

	template <typename T, std::size_t N, std::size_t M>
	vector<T, N> affine_transform<T, N, M>::translation_vector() const
	{
		vector<T, N> result;
		for (std::size_t r = 0; r < N; ++r)
			result[r] = m[r][M];
		return result;
	}

	template <typename T, std::size_t N, std::size_t M>
	matrix<T, N + 1, M + 1> affine_transform<T, N, M>::homogeneous_matrix() const
	{
		matrix<T, N + 1, M + 1> result;
		for (std::size_t r = 0; r < N; ++r)
		{
			for (std::size_t c = 0; c <= M; ++c)
				result[r][c] = m[r][c];
		}
		for (std::size_t c = 0; c < M; ++c)
			result[N][c] = T{0};
		result[N][M] = T{1};
		return result;
	}

	template <typename T, std::size_t N, std::size_t M>
	vector<T, N> affine_transform<T, N, M>::operator()(vector<T, M> const & v) const
	{
		vector<T, N> result = vector<T, N>::zero();
		for (std::size_t r = 0; r < N; ++r)
		{
			for (std::size_t c = 0; c < M; ++c)
				result[r] += m[r][c] * v[c];
		}
		return result;
	}

	template <typename T, std::size_t N, std::size_t M>
	point<T, N> affine_transform<T, N, M>::operator()(point<T, M> const & p) const
	{
		point<T, N> result = point<T, N>::zero();
		for (std::size_t r = 0; r < N; ++r)
		{
			for (std::size_t c = 0; c < M; ++c)
				result[r] += m[r][c] * p[c];
			result[r] += m[r][M];
		}
		return result;
	}

	template <typename T, std::size_t N, std::size_t M, std::size_t K>
	affine_transform<T, N, K> operator * (affine_transform<T, N, M> const & t1, affine_transform<T, M, K> const & t2)
	{
		auto const t1_lin = t1.linear_matrix();
		return {t1_lin * t2.linear_matrix(), t1_lin * t2.translation_vector() + t1.translation_vector()};
	}

	template <typename T, std::size_t N>
	std::optional<affine_transform<T, N, N>> inverse(affine_transform<T, N, N> const & t)
	{
		auto lin_inv = inverse(t.linear_matrix());
		if (!lin_inv)
			return std::nullopt;

		return affine_transform<T, N, N>{*lin_inv, - (*lin_inv) * t.translation_vector()};
	}

}
