#pragma once

#include <psemek/math/vector.hpp>
#include <psemek/math/point.hpp>
#include <psemek/math/quaternion.hpp>
#include <psemek/math/matrix.hpp>
#include <psemek/math/affine_transform.hpp>
#include <psemek/math/rotation.hpp>

#include <cstdint>

namespace psemek::gfx
{

	struct bone
	{
		static constexpr std::uint32_t null = static_cast<std::uint32_t>(-1);

		std::uint32_t parent = null;
		math::vector<float, 3> offset = math::vector<float, 3>::zero();
		math::matrix<float, 3, 3> axes = math::matrix<float, 3, 3>::identity();
	};

	template <typename T>
	struct bone_transform
	{
		math::quaternion<T> rotation;
		T scale;
		math::vector<T, 3> translation;

		static bone_transform identity();

		static bone_transform from_rotation(math::quaternion<T> const & rotation);
		static bone_transform from_scale(T const & scale);
		static bone_transform from_translation(math::vector<T, 3> const & translation);

		math::matrix<T, 3, 4> matrix() const;
	};

	template <typename T>
	bone_transform<T> bone_transform<T>::identity()
	{
		return {math::quaternion<T>::identity(), static_cast<T>(1), math::vector<T, 3>::zero()};
	}

	template <typename T>
	bone_transform<T> bone_transform<T>::from_rotation(math::quaternion<T> const & rotation)
	{
		return {rotation, static_cast<T>(1), math::vector<T, 3>::zero()};
	}

	template <typename T>
	bone_transform<T> bone_transform<T>::from_scale(T const & scale)
	{
		return {math::quaternion<T>::identity(), scale, math::vector<T, 3>::zero()};
	}

	template <typename T>
	bone_transform<T> bone_transform<T>::from_translation(math::vector<T, 3> const & translation)
	{
		return {math::quaternion<T>::identity(), static_cast<T>(1), translation};
	}

	template <typename T>
	math::matrix<T, 3, 4> bone_transform<T>::matrix() const
	{
		return math::affine_transform<T, 3, 3>(scale * math::quaternion_rotation<T>(math::quaternion<T>(rotation.coords)).linear_matrix(), translation).affine_matrix();
	}

	template <typename T>
	bone_transform<T> operator * (bone_transform<T> const & m1, bone_transform<T> const & m2)
	{
		// (1, t1) * (S1, 0) * (R1, 0) * (1, t2) * (S2, 0) * (R2, 0) =
		// (1, t1) * (S1, 0) * (1, R1 * t2) * (R1, 0) * (S2, 0) * (R2, 0) =
		// (1, t1) * (1, S1 * R1 * t2) * (S1, 0) * (R1, 0) * (S2, 0) * (R2, 0) =
		// (1, t1) * (1, S1 * R1 * t2) * (S1, 0) * (S2, 0) * (R1, 0) * (R2, 0) =
		// (1, t1 + S1 * R1 * t2) * (S1 * S2, 0) * (R1 * R2, 0)

		return bone_transform<T>{m1.rotation * m2.rotation, m1.scale * m2.scale, m1.translation + m1.scale * math::rotate(m1.rotation, m2.translation)};
	}

	template <typename T>
	bone_transform<T> inverse(bone_transform<T> const & m)
	{
		// [(1, T) * (S, 0) * (R, 0)]^-1 = (R^-1, 0) * (S^-1, 0) * (1, -T) =
		// = (1, - R^-1 S^-1 T) * (S^-1, 0) * (R^-1, 0)

		auto ir = math::inverse(m.rotation);
		return bone_transform<T>{ir, T{1} / m.scale, - math::rotate(ir, m.translation) / m.scale};
	}

	template <typename T>
	bone_transform<T> lerp(bone_transform<T> const & m1, bone_transform<T> const & m2, T t)
	{
		return {math::slerp(m1.rotation, m2.rotation, t), std::exp(math::lerp(std::log(m1.scale), std::log(m2.scale), t)), math::lerp(m1.translation, m2.translation, t)};
	}

	template <typename T>
	math::vector<T, 3> operator * (bone_transform<T> const & m, math::vector<T, 3> const & v)
	{
		return m.translation + m.scale * math::rotate(m.rotation, v);
	}

	template <typename T>
	math::point<T, 3> operator * (bone_transform<T> const & m, math::point<T, 3> const & p)
	{
		return p.zero() + m * (p - p.zero());
	}

	template <typename T>
	using pose = std::vector<bone_transform<T>>;

	template <typename T, typename Pose1, typename Pose2>
	void lerp(Pose1 const & pose1, Pose2 const & pose2, T t, pose<T> & output)
	{
		output.resize(pose1.size());
		for (std::size_t i = 0; i < output.size(); ++i)
			output[i] = lerp(pose1[i], pose2[i], t);
	}

	struct pose_compile_use_bone_axes_t {};
	static const pose_compile_use_bone_axes_t use_bone_axes;

	template <typename T, typename Pose, typename Bones>
	void compile(Pose const & local_pose, Bones const & bones, bone_transform<T> const & transform, pose<T> & result, pose_compile_use_bone_axes_t)
	{
		assert(local_pose.size() == bones.size());
		result.resize(local_pose.size());

		for (std::size_t b = 0; b < bones.size(); ++b)
		{
			result[b] =
				  gfx::bone_transform<float>{math::quaternion<float>::identity(), 1.f, bones[b].offset}
				* gfx::bone_transform<float>{math::quaternion<float>::rotation(bones[b].axes), 1.f, {0.f, 0.f, 0.f}}
				* local_pose[b]
				* gfx::bone_transform<float>{math::inverse(math::quaternion<float>::rotation(bones[b].axes)), 1.f, {0.f, 0.f, 0.f}}
				* gfx::bone_transform<float>{math::quaternion<float>::identity(), 1.f, -bones[b].offset}
				;

			auto p = bones[b].parent;
			if (p == gfx::bone::null)
				result[b] = transform * result[b];
			else
				result[b] = result[p] * result[b];
		}
	}

	template <typename T, typename Pose, typename Bones>
	void compile(Pose const & local_pose, Bones const & bones, bone_transform<T> const & transform, pose<T> & result)
	{
		assert(local_pose.size() == bones.size());
		result.resize(local_pose.size());

		for (std::size_t b = 0; b < bones.size(); ++b)
		{
			result[b] =
				  gfx::bone_transform<float>{math::quaternion<float>::identity(), 1.f, bones[b].offset}
				* local_pose[b]
				* gfx::bone_transform<float>{math::quaternion<float>::identity(), 1.f, -bones[b].offset}
				;

			auto p = bones[b].parent;
			if (p == gfx::bone::null)
				result[b] = transform * result[b];
			else
				result[b] = result[p] * result[b];
		}
	}

}
