#pragma once

#include <psemek/math/affine_transform.hpp>
#include <psemek/math/quaternion.hpp>

#include <psemek/util/assert.hpp>

namespace psemek::math
{

	// Rotation in oriented plane (i,j)
	template <typename T, std::size_t N>
	struct plane_rotation
	{
		std::size_t i, j;
		T angle;

		plane_rotation();
		plane_rotation(std::size_t i, std::size_t j, T angle);
		plane_rotation(plane_rotation const &) = default;

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

		affine_transform<T, N, N> transform() const;

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

	private:
		template <typename Matrix>
		void fill_matrix(Matrix & m) const;
	};

	// 3D-rotation around an axis
	template <typename T>
	struct axis_rotation
	{
		vector<T, 3> axis;
		T angle;

		axis_rotation();
		axis_rotation(vector<T, 3> const & axis, T angle);
		axis_rotation(axis_rotation const &) = default;

		matrix<T, 3, 4> affine_matrix() const;
		matrix<T, 3, 3> linear_matrix() const;
		vector<T, 3> translation_vector() const;
		matrix<T, 4, 4> homogeneous_matrix() const;

		affine_transform<T, 3, 3> transform() const;

		vector<T, 3> operator()(vector<T, 3> const & v) const;
		point<T, 3> operator()(point<T, 3> const & p) const;

	private:
		template <typename Matrix>
		void fill_matrix(Matrix & m) const;
	};

	template <typename T>
	struct quaternion_rotation
	{
		quaternion<T> quat;

		quaternion_rotation();
		quaternion_rotation(quaternion<T> const & quat);
		quaternion_rotation(T angle, vector<T, 3> const & axis);

		matrix<T, 3, 4> affine_matrix() const;
		matrix<T, 3, 3> linear_matrix() const;
		vector<T, 3> translation_vector() const;
		matrix<T, 4, 4> homogeneous_matrix() const;

		affine_transform<T, 3, 3> transform() const;

		vector<T, 3> operator()(vector<T, 3> const & v) const;
		point<T, 3> operator()(point<T, 3> const & p) const;

	private:
		template <typename Matrix>
		void fill_matrix(Matrix & m) const;
	};

	template <typename T>
	axis_rotation(vector<T, 3>, T) -> axis_rotation<T>;

	template <typename T, std::size_t N>
	plane_rotation<T, N>::plane_rotation()
		: i{0}
		, j{1}
		, angle{0}
	{
		assert(i < N);
		assert(j < N);
	}

	template <typename T, std::size_t N>
	plane_rotation<T, N>::plane_rotation(std::size_t i, std::size_t j, T angle)
		: i{i}
		, j{j}
		, angle{angle}
	{
		assert(i < N);
		assert(j < N);
	}

	template <typename T, std::size_t N>
	matrix<T, N, N + 1> plane_rotation<T, N>::affine_matrix() const
	{
		auto result = matrix<T, N, N + 1>::identity();
		fill_matrix(result);
		return result;
	}

	template <typename T, std::size_t N>
	matrix<T, N, N> plane_rotation<T, N>::linear_matrix() const
	{
		auto result = matrix<T, N, N>::identity();
		fill_matrix(result);
		return result;
	}

	template <typename T, std::size_t N>
	vector<T, N> plane_rotation<T, N>::translation_vector() const
	{
		return vector<T, N>::zero();
	}

	template <typename T, std::size_t N>
	matrix<T, N + 1, N + 1> plane_rotation<T, N>::homogeneous_matrix() const
	{
		auto result = matrix<T, N + 1, N + 1>::identity();
		fill_matrix(result);
		return result;
	}

	template <typename T, std::size_t N>
	affine_transform<T, N, N> plane_rotation<T, N>::transform() const
	{
		return {affine_matrix()};
	}

	template <typename T, std::size_t N>
	vector<T, N> plane_rotation<T, N>::operator()(vector<T, N> const & v) const
	{
		auto result = v;
		result[i] = v[i] * std::cos(angle) - v[j] * std::sin(angle);
		result[j] = v[i] * std::sin(angle) + v[j] * std::cos(angle);
		return result;
	}

	template <typename T, std::size_t N>
	point<T, N> plane_rotation<T, N>::operator()(point<T, N> const & p) const
	{
		auto result = p;
		result[i] = p[i] * std::cos(angle) - p[j] * std::sin(angle);
		result[j] = p[i] * std::sin(angle) + p[j] * std::cos(angle);
		return result;
	}

	template <typename T, std::size_t N>
	template <typename Matrix>
	void plane_rotation<T, N>::fill_matrix(Matrix & m) const
	{
		m[i][i] =  std::cos(angle);
		m[i][j] = -std::sin(angle);
		m[j][i] =  std::sin(angle);
		m[j][j] =  std::cos(angle);
	}

	template <typename T, std::size_t N>
	plane_rotation<T, N> inverse(plane_rotation<T, N> const & r)
	{
		return {r.j, r.i, r.angle};
	}

	template <typename T>
	axis_rotation<T>::axis_rotation()
		: axis{T{0}, T{0}, T{1}}
		, angle{T{0}}
	{}

	template <typename T>
	axis_rotation<T>::axis_rotation(vector<T, 3> const & axis, T angle)
		: axis{axis}
		, angle{angle}
	{}

	template <typename T>
	matrix<T, 3, 4> axis_rotation<T>::affine_matrix() const
	{
		auto result = matrix<T, 3, 4>::identity();
		fill_matrix(result);
		return result;
	}

	template <typename T>
	matrix<T, 3, 3> axis_rotation<T>::linear_matrix() const
	{
		auto result = matrix<T, 3, 3>::identity();
		fill_matrix(result);
		return result;
	}

	template <typename T>
	vector<T, 3> axis_rotation<T>::translation_vector() const
	{
		return vector<T, 3>::zero();
	}

	template <typename T>
	matrix<T, 4, 4> axis_rotation<T>::homogeneous_matrix() const
	{
		auto result = matrix<T, 4, 4>::identity();
		fill_matrix(result);
		return result;
	}

	template <typename T>
	affine_transform<T, 3, 3> axis_rotation<T>::transform() const
	{
		return {affine_matrix()};
	}

	template <typename T>
	vector<T, 3> axis_rotation<T>::operator()(vector<T, 3> const & v) const
	{
		return linear_matrix() * v;
	}

	template <typename T>
	point<T, 3> axis_rotation<T>::operator()(point<T, 3> const & p) const
	{
		auto const o = math::point<T, 3>::zero();
		return linear_matrix() * (p - o) + o;
	}

	template <typename T>
	template <typename Matrix>
	void axis_rotation<T>::fill_matrix(Matrix & m) const
	{
		T const c = std::cos(angle);
		T const s = std::sin(angle);
		T const x = axis[0];
		T const y = axis[1];
		T const z = axis[2];
		m[0][0] = c + x * x * (1 - c);
		m[0][1] = x * y * (1 - c) - z * s;
		m[0][2] = x * z * (1 - c) + y * s;
		m[1][0] = y * x * (1 - c) + z * s;
		m[1][1] = c + y * y * (1 - c);
		m[1][2] = y * z * (1 - c) - x * s;
		m[2][0] = z * x * (1 - c) - y * s;
		m[2][1] = z * y * (1 - c) + x * s;
		m[2][2] = c + z * z * (1 - c);
	}

	template <typename T>
	axis_rotation<T> inverse(axis_rotation<T> const & r)
	{
		return {r.axis, -r.angle};
	}

	template <typename T>
	quaternion_rotation<T>::quaternion_rotation()
		: quat{quaternion<T>::identity()}
	{}

	template <typename T>
	quaternion_rotation<T>::quaternion_rotation(quaternion<T> const & quat)
		: quat{quat}
	{}

	template <typename T>
	quaternion_rotation<T>::quaternion_rotation(T angle, vector<T, 3> const & axis)
		: quat{quaternion<T>::rotation(angle, axis)}
	{}

	template <typename T>
	matrix<T, 3, 4> quaternion_rotation<T>::affine_matrix() const
	{
		auto result = matrix<T, 3, 4>::identity();
		fill_matrix(result);
		return result;
	}

	template <typename T>
	matrix<T, 3, 3> quaternion_rotation<T>::linear_matrix() const
	{
		auto result = matrix<T, 3, 3>::identity();
		fill_matrix(result);
		return result;
	}

	template <typename T>
	vector<T, 3> quaternion_rotation<T>::translation_vector() const
	{
		return vector<T, 3>::zero();
	}

	template <typename T>
	matrix<T, 4, 4> quaternion_rotation<T>::homogeneous_matrix() const
	{
		auto result = matrix<T, 4, 4>::identity();
		fill_matrix(result);
		return result;
	}

	template <typename T>
	affine_transform<T, 3, 3> quaternion_rotation<T>::transform() const
	{
		return {affine_matrix()};
	}

	template <typename T>
	vector<T, 3> quaternion_rotation<T>::operator()(vector<T, 3> const & v) const
	{
		return linear_matrix() * v;
	}

	template <typename T>
	point<T, 3> quaternion_rotation<T>::operator()(point<T, 3> const & p) const
	{
		static constexpr auto o = point<T, 3>::zero();
		return o + linear_matrix() * (p - o);
	}

	template <typename T>
	template <typename Matrix>
	void quaternion_rotation<T>::fill_matrix(Matrix & m) const
	{
		m[0][0] = 1 - 2 * (quat[1] * quat[1] + quat[2] * quat[2]);
		m[0][1] = 2 * (quat[0] * quat[1] - quat[3] * quat[2]);
		m[0][2] = 2 * (quat[0] * quat[2] + quat[3] * quat[1]);
		m[1][0] = 2 * (quat[0] * quat[1] + quat[3] * quat[2]);
		m[1][1] = 1 - 2 * (quat[0] * quat[0] + quat[2] * quat[2]);
		m[1][2] = 2 * (quat[1] * quat[2] - quat[3] * quat[0]);
		m[2][0] = 2 * (quat[0] * quat[2] - quat[3] * quat[1]);
		m[2][1] = 2 * (quat[1] * quat[2] + quat[3] * quat[0]);
		m[2][2] = 1 - 2 * (quat[0] * quat[0] + quat[1] * quat[1]);
	}

	template <typename T>
	quaternion_rotation<T> inverse(quaternion_rotation<T> const & r)
	{
		return quaternion_rotation<T>{quaternion<T>{-r.quat[0], -r.quat[1], -r.quat[2], r.quat[3]}};
	}

}
