#pragma once

#include <psemek/math/vector.hpp>
#include <psemek/math/math.hpp>
#include <psemek/math/matrix.hpp>
#include <psemek/math/interval.hpp>

#include <format>

namespace psemek::math
{

	template <typename T>
	struct quaternion
	{
		math::vector<T, 4> coords;

		quaternion() = default;

		static quaternion<T> zero();
		static quaternion<T> identity();
		static quaternion<T> i();
		static quaternion<T> j();
		static quaternion<T> k();
		static quaternion<T> scalar(T value);
		static quaternion<T> vector(math::vector<T, 3> const & v);
		static quaternion<T> rotation(T angle, math::vector<T, 3> const & axis);
		static quaternion<T> rotation(matrix<T, 3, 3> const & m);

		quaternion(math::vector<T, 4> const & v)
			: coords{v}
		{}

		T & operator[] (std::size_t i) { return coords[i]; }
		T const & operator[] (std::size_t i) const { return coords[i]; }

		quaternion<T> & operator += (quaternion<T> const & other);
		quaternion<T> & operator -= (quaternion<T> const & other);
	};

	template <typename ... Args>
	quaternion(Args && ...) -> quaternion<std::common_type_t<Args...>>;

	template <typename T>
	quaternion<T> quaternion<T>::zero()
	{
		return quaternion<T>{math::vector<T, 4>::zero()};
	}

	template <typename T>
	quaternion<T> quaternion<T>::identity()
	{
		return quaternion<T>{{T(0), T(0), T(0), T(1)}};
	}

	template <typename T>
	quaternion<T> quaternion<T>::i()
	{
		return  quaternion<T>{{T(1), T(0), T(0), T(0)}};
	}

	template <typename T>
	quaternion<T> quaternion<T>::j()
	{
		return  quaternion<T>{{T(0), T(1), T(0), T(0)}};
	}

	template <typename T>
	quaternion<T> quaternion<T>::k()
	{
		return  quaternion<T>{{T(0), T(0), T(1), T(0)}};
	}

	template <typename T>
	quaternion<T> quaternion<T>::scalar(T value)
	{
		return {{T(0), T(0), T(0), value}};
	}

	template <typename T>
	quaternion<T> quaternion<T>::vector(math::vector<T, 3> const & v)
	{
		return {{v[0], v[1], v[2], T(0)}};
	}

	template <typename T>
	quaternion<T> quaternion<T>::rotation(T angle, math::vector<T, 3> const & axis)
	{
		auto const c = std::cos(angle / 2);
		auto const s = std::sin(angle / 2);

		return quaternion<T>{{s * axis[0], s * axis[1], s * axis[2], c}};
	}

	template <typename T>
	quaternion<T> quaternion<T>::rotation(matrix<T, 3, 3> const & m)
	{
		// https://www.euclideanspace.com/maths/geometry/rotations/conversions/matrixToQuaternion/

		auto tr = trace(m);

		if (tr > 0)
		{
			auto S = std::sqrt(tr + 1) * 2;
			return {{(m[2][1] - m[1][2]) / S, (m[0][2] - m[2][0]) / S, (m[1][0] - m[0][1]) / S, S / 4}};
		}
		else if (m[0][0] > m[1][1] && m[0][0] > m[2][2])
		{
			auto S = std::sqrt(1 + m[0][0] - m[1][1] - m[2][2]) * 2;
			return {{S / 4, (m[0][1] + m[1][0]) / S, (m[0][2] + m[2][0]) / S, (m[2][1] - m[1][2]) / S}};
		}
		else if (m[1][1] > m[2][2])
		{
			auto S = std::sqrt(1 - m[0][0] + m[1][1] - m[2][2]) * 2;
			return {{(m[0][1] + m[1][0]) / S, S / 4, (m[1][2] + m[2][1]) / S, (m[0][2] - m[2][0]) / S}};
		}
		else
		{
			auto S = std::sqrt(1 - m[0][0] - m[1][1] + m[2][2]) * 2;
			return {{(m[0][2] + m[2][0]) / S, (m[1][2] + m[2][1]) / S, S / 4, (m[1][0] - m[0][1]) / S}};
		}
	}

	template <typename T>
	quaternion<T> operator - (quaternion<T> const & q)
	{
		return {-q.coords};
	}

	template <typename T>
	quaternion<T> operator + (quaternion<T> const & q1, quaternion<T> const & q2)
	{
		return {q1.coords + q2.coords};
	}

	template <typename T>
	quaternion<T> operator - (quaternion<T> const & q1, quaternion<T> const & q2)
	{
		return {q1.coords - q2.coords};
	}

	template <typename T>
	quaternion<T> operator * (quaternion<T> const & q, T s)
	{
		return {q.coords * s};
	}

	template <typename T>
	quaternion<T> operator * (T s, quaternion<T> const & q)
	{
		return {s * q.coords};
	}

	template <typename T>
	quaternion<T> operator / (quaternion<T> const & q, T s)
	{
		return {q.coords / s};
	}

	template <typename T>
	quaternion<T> operator * (quaternion<T> const & q1, quaternion<T> const & q2)
	{
		quaternion<T> r;
		r[0] = q1[3] * q2[0] + q1[0] * q2[3] + q1[1] * q2[2] - q1[2] * q2[1];
		r[1] = q1[3] * q2[1] - q1[0] * q2[2] + q1[1] * q2[3] + q1[2] * q2[0];
		r[2] = q1[3] * q2[2] + q1[0] * q2[1] - q1[1] * q2[0] + q1[2] * q2[3];
		r[3] = q1[3] * q2[3] - q1[0] * q2[0] - q1[1] * q2[1] - q1[2] * q2[2];
		return r;
	}

	template <typename T>
	quaternion<T> & quaternion<T>::operator += (quaternion<T> const & other)
	{
		(*this) = (*this) + other;
		return *this;
	}

	template <typename T>
	quaternion<T> & quaternion<T>::operator -= (quaternion<T> const & other)
	{
		(*this) = (*this) - other;
		return *this;
	}

	template <typename T>
	T length_sqr(quaternion<T> const & q)
	{
		return length_sqr(q.coords);
	}

	template <typename T>
	T length(quaternion<T> const & q)
	{
		return length(q.coords);
	}

	template <typename T>
	quaternion<T> normalized(quaternion<T> const & q)
	{
		return {normalized(q.coords)};
	}

	template <typename T>
	quaternion<T> conjugate(quaternion<T> const & q)
	{
		return {{-q[0], -q[1], -q[2], q[3]}};
	}

	template <typename T>
	quaternion<T> inverse(quaternion<T> const & q)
	{
		return conjugate(q) / length(q);
	}

	template <typename T>
	vector<T, 3> rotate(quaternion<T> const & q, vector<T, 3> const & v)
	{
		// 2x faster than q*v*cong(q), see https://www.johndcook.com/blog/2021/06/16/faster-quaternion-rotations/
		auto qv = vector{q[0], q[1], q[2]};
		auto t = T(2) * cross(qv, v);
		return v + q[3] * t + cross(qv, t);
	}

	template <typename T>
	vector<T, 3> rotate(quaternion<T> const & q1, vector<T, 3> const & v, quaternion<T> const & q2)
	{
		auto res = q1 * quaternion<T>::vector(v) * conjugate(q2);
		return {res[0], res[1], res[2]};
	}

	// v0 & v1 assumed to be normalized
	template <typename T>
	quaternion<T> shortest_arc(vector<T, 3> const & v0, vector<T, 3> const & v1, T const eps = T{1e-6})
	{
		auto axis = cross(v0, v1);
		auto d = dot(v0, v1);

		if (d + T{1} < eps)
			axis = ort(v0);

		quaternion<T> result;

		result[0] = axis[0];
		result[1] = axis[1];
		result[2] = axis[2];
		result[3] = T{1} + d;

		result.coords = normalized(result.coords);
		return result;
	}

	template <typename T>
	T angle(quaternion<T> const & q0, quaternion<T> const & q1)
	{
		return std::acos(std::min(T{1}, std::abs(dot(q0.coords, q1.coords))));
	}

	template <typename T>
	quaternion<T> slerp(quaternion<T> const & q0, quaternion<T> const & q1, T const & t)
	{
		using std::sin;
		using std::acos;
		using std::abs;

		auto const d = dot(q0.coords, q1.coords);

		// prevent division by zero
		if (d >= T{1})
			return quaternion<T>{lerp(q0.coords, q1.coords, t)};

		auto const omega = acos(abs(d));

		// NB: the case of omega ~ pi is ambiguous and isn't handled in any special way

		auto const s = sin(omega);
		auto const w0 = sin((1 - t) * omega) / s;
		auto const w1 = sin(t * omega) / s * ((d > T{0}) ? T{1} : -T{1});
		return quaternion<T>{q0.coords * w0 + q1.coords * w1};
	}

	template <typename T>
	quaternion<T> exp(quaternion<T> const & q)
	{
		auto v = vector{q[0], q[1], q[2]};
		auto l = length(v);

		auto s = sin_over_x(l);

		return std::exp(q[3]) * quaternion<T>({v[0] * s, v[1] * s, v[2] * s, std::cos(l)});
	}

	template <typename T>
	quaternion<T> log(quaternion<T> const & q)
	{
		auto v = vector{q[0], q[1], q[2]};
		auto l = length(v);
		auto n = length(q);

		if (l == T{})
			v = v.zero();
		else
			v = (v / l) * std::acos(q[3] / n);

		return quaternion<T>{{v[0], v[1], v[2], std::log(n)}};
	}

	template <typename T, typename H>
	quaternion<T> cast(quaternion<H> const & q)
	{
		return {cast<T>(q.coords)};
	}

	template <typename T>
	std::ostream & operator << (std::ostream & os, quaternion<T> const & q)
	{
		return os << q.coords;
	}

	// Returns a matrix M such that the angular gradient of function f(q) is M * grad(f)
	template <typename T>
	math::matrix<T, 3, 4> angular_gradient(quaternion<T> const & q)
	{
		math::matrix<T, 3, 4> result;
		result[0][0] =   q[3];
		result[0][1] = - q[2];
		result[0][2] =   q[1];
		result[0][3] = - q[0];
		result[1][0] =   q[2];
		result[1][1] =   q[3];
		result[1][2] = - q[0];
		result[1][3] = - q[1];
		result[2][0] = - q[1];
		result[2][1] =   q[0];
		result[2][2] =   q[3];
		result[2][3] = - q[2];
		return result;
	}

}

namespace std
{

	template <typename T, typename Char>
	struct formatter<::psemek::math::quaternion<T>, Char>
		: formatter<::psemek::math::vector<T, 4>>
	{
		using formatter<::psemek::math::vector<T, 4>>::parse;

		template <typename FormatContext>
		auto format(::psemek::math::quaternion<T> const & q, FormatContext & ctx) const
		{
			formatter<::psemek::math::vector<T, 4>>::format(q.coords, ctx);
			return ctx.out();
		}
	};

}
