/* +------------------------------------------------------------------------+
   |                     Mobile Robot Programming Toolkit (MRPT)            |
   |                          https://www.mrpt.org/                         |
   |                                                                        |
   | Copyright (c) 2005-2022, Individual contributors, see AUTHORS file     |
   | See: https://www.mrpt.org/Authors - All rights reserved.               |
   | Released under BSD License. See: https://www.mrpt.org/License          |
   +------------------------------------------------------------------------+ */

#pragma once

#include <mrpt/core/bits_math.h>  // square()
#include <mrpt/core/exceptions.h>
#include <mrpt/math/CMatrixDynamic.h>
#include <mrpt/math/CVectorFixed.h>

namespace mrpt::math
{
// For use with a constructor of CQuaternion
enum TConstructorFlags_Quaternions
{
	UNINITIALIZED_QUATERNION = 0
};

/** A quaternion, which can represent a 3D rotation as pair \f$ (r,\mathbf{u})
 *\f$, with a real part "r" and a 3D vector \f$ \mathbf{u} = (x,y,z) \f$, or
 *alternatively, q = r + ix + jy + kz.
 *
 *  The elements of the quaternion can be accessed by either:
 *		- r()(equivalent to w()), x(), y(), z(), or
 *		- the operator [] with indices: [0]=`w`, [1]:`x`, [2]:`y`, [3]:`z`
 *
 *  Users will usually employ the type `CQuaternionDouble` instead of this
 *template.
 *
 * For more information about quaternions, see:
 *  - http://people.csail.mit.edu/bkph/articles/Quaternions.pdf
 *  - http://en.wikipedia.org/wiki/Quaternions_and_spatial_rotation
 *
 * \ingroup mrpt_math_grp
 * \sa mrpt::poses::CPose3D
 */
template <class T>
class CQuaternion : public CVectorFixed<T, 4>
{
	using Base = CVectorFixed<T, 4>;

   public:
	/* @{ Constructors
	 */

	/**	Can be used with UNINITIALIZED_QUATERNION as argument, does not
	 * initialize the 4 elements of the quaternion (use this constructor when
	 * speed is critical). */
	inline CQuaternion(TConstructorFlags_Quaternions) {}
	/**	Default constructor: construct a (1, (0,0,0) ) quaternion representing
	 * no rotation. */
	inline CQuaternion()
	{
		r() = 1;
		x() = 0;
		y() = 0;
		z() = 0;
	}

	/**	Construct a quaternion from its parameters 'r', 'x', 'y', 'z', with q =
	 * r + ix + jy + kz. */
	inline CQuaternion(const T R, const T X, const T Y, const T Z)
	{
		r() = R;
		x() = X;
		y() = Y;
		z() = Z;
		ensurePositiveRealPart();

		ASSERTMSG_(
			std::abs(normSqr() - 1.0) < 1e-3,
			mrpt::format(
				"Initialization data for quaternion is not normalized: %f %f "
				"%f %f -> sqrNorm=%f",
				R, X, Y, Z, normSqr()));
	}

	/* @}
	 */

	/** Adhere to the convention of w>=0 to avoid ambiguity of quaternion double
	 * cover of SO(3) */
	inline void ensurePositiveRealPart()
	{
		// Ensure r()>0
		if (r() < 0)
		{
			r() = -r();
			x() = -x();
			y() = -y();
			z() = -z();
		}
	}

	/** Return r (real part) coordinate of the quaternion */
	inline T r() const { return (*this)[0]; }
	/** Return w (real part) coordinate of the quaternion. Alias of r() */
	inline T w() const { return (*this)[0]; }
	/** Return x coordinate of the quaternion */
	inline T x() const { return (*this)[1]; }
	/** Return y coordinate of the quaternion */
	inline T y() const { return (*this)[2]; }
	/** Return z coordinate of the quaternion */
	inline T z() const { return (*this)[3]; }
	/** Set r (real part) coordinate of the quaternion */
	inline void r(const T r) { (*this)[0] = r; }
	/** Set w (real part) coordinate of the quaternion. Alias of r() */
	inline void w(const T w) { (*this)[0] = w; }
	/** Set x coordinate of the quaternion */
	inline void x(const T x) { (*this)[1] = x; }
	/** Set y coordinate of the quaternion */
	inline void y(const T y) { (*this)[2] = y; }
	/** Set z coordinate of the quaternion */
	inline void z(const T z) { (*this)[3] = z; }

	inline T& r() { return (*this)[0]; }
	inline T& x() { return (*this)[1]; }
	inline T& y() { return (*this)[2]; }
	inline T& z() { return (*this)[3]; }

	/**	Set this quaternion to the rotation described by a 3D (Rodrigues)
	 * rotation vector \f$ \mathbf{v} \f$:
	 *   If \f$ \mathbf{v}=0 \f$, then the quaternion is \f$ \mathbf{q} = [1 ~
	 * 0 ~ 0 ~ 0]^\top \f$, otherwise:
	 *    \f[ \mathbf{q} = \left[ \begin{array}{c}
	 *       \cos(\frac{\theta}{2}) \\
	 *       v_x \frac{\sin(\frac{\theta}{2})}{\theta} \\
	 *       v_y \frac{\sin(\frac{\theta}{2})}{\theta} \\
	 *       v_z \frac{\sin(\frac{\theta}{2})}{\theta}
	 *    \end{array} \right] \f]
	 *    where \f$ \theta = |\mathbf{v}| = \sqrt{v_x^2+v_y^2+v_z^2} \f$.
	 *  \sa "Representing Attitude: Euler Angles, Unit Quaternions, and
	 * Rotation Vectors (2006)", James Diebel.
	 */
	template <class ARRAYLIKE3>
	void fromRodriguesVector(const ARRAYLIKE3& v)
	{
		MRPT_START
		ASSERT_(v.size() == 3);
		const T x = v[0], y = v[1], z = v[2];
		const T theta_sq = x * x + y * y + z * z, theta = std::sqrt(theta_sq);
		T r, i;
		if (theta < 1e-6)
		{
			// Taylor series approximation:
			const T theta_po4 = theta_sq * theta_sq;
			i = T(0.5) - T(1.0 / 48.0) * theta_sq + T(1.0 / 3840.0) * theta_po4;
			r = T(1.0) - T(0.5) * theta_sq + T(1.0 / 384.0) * theta_po4;
		}
		else
		{
			i = (std::sin(theta / 2)) / theta;
			r = std::cos(theta / 2);
		}
		(*this)[0] = r;
		(*this)[1] = x * i;
		(*this)[2] = y * i;
		(*this)[3] = z * i;
		ASSERTMSG_(
			normSqr() - 1.0 < 1e-6,
			mrpt::format(
				"fromRodriguesVector() failed, tangent_vector=[%g %g %g]", v[0],
				v[1], v[2]));
		MRPT_END
	}

	/** @name Lie Algebra methods
		@{ */

	/** Logarithm of the 3x3 matrix defined by this pose, generating the
	 * corresponding vector in the SO(3) Lie Algebra,
	 *  which coincides with the so-called "rotation vector" (I don't have
	 * space here for the proof ;-).
	 *  \param[out] out_ln The target vector, which can be: std::vector<>, or
	 * mrpt::math::CVectorDouble or any row or column Eigen::Matrix<>.
	 *  \sa exp,  mrpt::poses::SE_traits  */
	template <class ARRAYLIKE3>
	inline void ln(ARRAYLIKE3& out_ln) const
	{
		if (out_ln.size() != 3) out_ln.resize(3);
		this->ln_noresize(out_ln);
	}
	/** overload that returns by value */
	template <class ARRAYLIKE3>
	inline ARRAYLIKE3 ln() const
	{
		ARRAYLIKE3 out_ln;
		this->ln(out_ln);
		return out_ln;
	}
	/** Like ln() but does not try to resize the output vector. */
	template <class ARRAYLIKE3>
	void ln_noresize(ARRAYLIKE3& out_ln) const
	{
		using mrpt::square;
		const T xyz_norm = std::sqrt(square(x()) + square(y()) + square(z()));
		const T K = (xyz_norm < 1e-7) ? 2 : 2 * ::acos(r()) / xyz_norm;
		out_ln[0] = K * x();
		out_ln[1] = K * y();
		out_ln[2] = K * z();
	}

	/** Exponential map from the SO(3) Lie Algebra to unit quaternions.
	 *  \sa ln,  mrpt::poses::SE_traits  */
	template <class ARRAYLIKE3>
	inline static CQuaternion<T> exp(const ARRAYLIKE3& v)
	{
		CQuaternion<T> q(UNINITIALIZED_QUATERNION);
		q.fromRodriguesVector(v);
		return q;
	}
	/** \overload */
	template <class ARRAYLIKE3>
	inline static void exp(const ARRAYLIKE3& v, CQuaternion<T>& out_quat)
	{
		out_quat.fromRodriguesVector(v);
	}

	/** @} */  // end of Lie algebra

	/**	Calculate the "cross" product (or "composed rotation") of two
	 * quaternion: this = q1 x q2
	 *   After the operation, "this" will represent the composed rotations of
	 * q1 and q2 (q2 applied "after" q1).
	 */
	inline void crossProduct(const CQuaternion& q1, const CQuaternion& q2)
	{
		// First: compute result, then save in this object. In this way we avoid
		// problems when q1 or q2 == *this !!
		const T new_r = q1.r() * q2.r() - q1.x() * q2.x() - q1.y() * q2.y() -
			q1.z() * q2.z();
		const T new_x = q1.r() * q2.x() + q2.r() * q1.x() + q1.y() * q2.z() -
			q2.y() * q1.z();
		const T new_y = q1.r() * q2.y() + q2.r() * q1.y() + q1.z() * q2.x() -
			q2.z() * q1.x();
		const T new_z = q1.r() * q2.z() + q2.r() * q1.z() + q1.x() * q2.y() -
			q2.x() * q1.y();
		this->r(new_r);
		this->x(new_x);
		this->y(new_y);
		this->z(new_z);
		this->normalize();
	}

	/** Rotate a 3D point (lx,ly,lz) -> (gx,gy,gz) as described by this
	 * quaternion
	 */
	void rotatePoint(
		const double lx, const double ly, const double lz, double& gx,
		double& gy, double& gz) const
	{
		const double t2 = r() * x();
		const double t3 = r() * y();
		const double t4 = r() * z();
		const double t5 = -x() * x();
		const double t6 = x() * y();
		const double t7 = x() * z();
		const double t8 = -y() * y();
		const double t9 = y() * z();
		const double t10 = -z() * z();
		gx = 2 * ((t8 + t10) * lx + (t6 - t4) * ly + (t3 + t7) * lz) + lx;
		gy = 2 * ((t4 + t6) * lx + (t5 + t10) * ly + (t9 - t2) * lz) + ly;
		gz = 2 * ((t7 - t3) * lx + (t2 + t9) * ly + (t5 + t8) * lz) + lz;
	}

	/** Rotate a 3D point (lx,ly,lz) -> (gx,gy,gz) as described by the inverse
	 * (conjugate) of this quaternion
	 */
	void inverseRotatePoint(
		const double lx, const double ly, const double lz, double& gx,
		double& gy, double& gz) const
	{
		const double t2 = -r() * x();
		const double t3 = -r() * y();
		const double t4 = -r() * z();
		const double t5 = -x() * x();
		const double t6 = x() * y();
		const double t7 = x() * z();
		const double t8 = -y() * y();
		const double t9 = y() * z();
		const double t10 = -z() * z();
		gx = 2 * ((t8 + t10) * lx + (t6 - t4) * ly + (t3 + t7) * lz) + lx;
		gy = 2 * ((t4 + t6) * lx + (t5 + t10) * ly + (t9 - t2) * lz) + ly;
		gz = 2 * ((t7 - t3) * lx + (t2 + t9) * ly + (t5 + t8) * lz) + lz;
	}

	/** Return the squared norm of the quaternion */
	inline double normSqr() const
	{
		using mrpt::square;
		return square(r()) + square(x()) + square(y()) + square(z());
	}

	/**	Normalize this quaternion, so its norm becomes the unitity.
	 */
	inline void normalize()
	{
		ensurePositiveRealPart();
		const T qq = 1.0 / std::sqrt(normSqr());
		for (unsigned int i = 0; i < 4; i++)
			(*this)[i] *= qq;
	}

	/** Calculate the 4x4 Jacobian of the normalization operation of this
	 * quaternion.
	 *  The output matrix can be a dynamic or fixed size (4x4) matrix.
	 */
	template <class MATRIXLIKE>
	void normalizationJacobian(MATRIXLIKE& J) const
	{
		const T n = 1.0 / std::pow(normSqr(), T(1.5));
		J.setSize(4, 4);
		J(0, 0) = x() * x() + y() * y() + z() * z();
		J(0, 1) = -r() * x();
		J(0, 2) = -r() * y();
		J(0, 3) = -r() * z();

		J(1, 0) = -x() * r();
		J(1, 1) = r() * r() + y() * y() + z() * z();
		J(1, 2) = -x() * y();
		J(1, 3) = -x() * z();

		J(2, 0) = -y() * r();
		J(2, 1) = -y() * x();
		J(2, 2) = r() * r() + x() * x() + z() * z();
		J(2, 3) = -y() * z();

		J(3, 0) = -z() * r();
		J(3, 1) = -z() * x();
		J(3, 2) = -z() * y();
		J(3, 3) = r() * r() + x() * x() + y() * y();
		J *= n;
	}

	/** Compute the Jacobian of the rotation composition operation \f$ p =
	 * f(\cdot) = q_{this} \times r \f$, that is the 4x4 matrix \f$
	 * \frac{\partial f}{\partial q_{this} }  \f$.
	 *  The output matrix can be a dynamic or fixed size (4x4) matrix.
	 */
	template <class MATRIXLIKE>
	inline void rotationJacobian(MATRIXLIKE& J) const
	{
		J.setSize(4, 4);
		J(0, 0) = r();
		J(0, 1) = -x();
		J(0, 2) = -y();
		J(0, 3) = -z();
		J(1, 0) = x();
		J(1, 1) = r();
		J(1, 2) = -z();
		J(1, 3) = y();
		J(2, 0) = y();
		J(2, 1) = z();
		J(2, 2) = r();
		J(2, 3) = -x();
		J(3, 0) = z();
		J(3, 1) = -y();
		J(3, 2) = x();
		J(3, 3) = r();
	}

	/** Calculate the 3x3 rotation matrix associated to this quaternion: \f[
	 * \mathbf{R} = \left(
	 *    \begin{array}{ccc}
	 *     q_r^2+q_x^2-q_y^2-q_z^2 	&  2(q_x q_y - q_r q_z)	&  	2(q_z
	 * q_x+q_r q_y) \\
	 *     2(q_x q_y+q_r q_z) 		& q_r^2-q_x^2+q_y^2-q_z^2 	& 2(q_y
	 * q_z-q_r q_x) \\
	 *     2(q_z q_x-q_r q_y) & 2(q_y q_z+q_r q_x)  & q_r^2- q_x^2 - q_y^2 +
	 * q_z^2
	 *    \end{array}
	 *  \right)\f]
	 *
	 */
	template <class MATRIXLIKE>
	inline void rotationMatrix(MATRIXLIKE& M) const
	{
		M.setSize(3, 3);
		rotationMatrixNoResize(M);
	}

	template <class MATRIXLIKE>
	inline MATRIXLIKE rotationMatrix() const
	{
		MATRIXLIKE M(3, 3);
		rotationMatrixNoResize(M);
		return M;
	}

	/** Fill out the top-left 3x3 block of the given matrix with the rotation
	 * matrix associated to this quaternion (does not resize the matrix, for
	 * that, see rotationMatrix). */
	template <class MATRIXLIKE>
	inline void rotationMatrixNoResize(MATRIXLIKE& M) const
	{
		M(0, 0) = r() * r() + x() * x() - y() * y() - z() * z();
		M(0, 1) = 2 * (x() * y() - r() * z());
		M(0, 2) = 2 * (z() * x() + r() * y());
		M(1, 0) = 2 * (x() * y() + r() * z());
		M(1, 1) = r() * r() - x() * x() + y() * y() - z() * z();
		M(1, 2) = 2 * (y() * z() - r() * x());
		M(2, 0) = 2 * (z() * x() - r() * y());
		M(2, 1) = 2 * (y() * z() + r() * x());
		M(2, 2) = r() * r() - x() * x() - y() * y() + z() * z();
	}

	/**	Return the conjugate quaternion  */
	inline void conj(CQuaternion& q_out) const
	{
		q_out.r(r());
		q_out.x(-x());
		q_out.y(-y());
		q_out.z(-z());
	}

	/**	Return the conjugate quaternion  */
	inline CQuaternion conj() const
	{
		CQuaternion q_aux;
		conj(q_aux);
		return q_aux;
	}

	/**	Return the yaw, pitch & roll angles associated to quaternion
	 *  \sa For the equations, see The MRPT Book, or see
	 * http://www.euclideanspace.com/maths/geometry/rotations/conversions/quaternionToEuler/Quaternions.pdf
	 *  \sa rpy_and_jacobian
	 */
	inline void rpy(T& roll, T& pitch, T& yaw) const
	{
		rpy_and_jacobian(
			roll, pitch, yaw, static_cast<mrpt::math::CMatrixDouble*>(nullptr));
	}

	/**	Return the yaw, pitch & roll angles associated to quaternion, and
	 * (optionally) the 3x4 Jacobian of the transformation.
	 *  Note that both the angles and the Jacobian have one set of normal
	 * equations, plus other special formulas for the degenerated cases of
	 * |pitch|=90 degrees.
	 *  \sa For the equations, see The MRPT Book, or
	 * http://www.euclideanspace.com/maths/geometry/rotations/conversions/quaternionToEuler/Quaternions.pdf
	 * \sa rpy
	 */
	template <class MATRIXLIKE>
	void rpy_and_jacobian(
		T& roll, T& pitch, T& yaw, MATRIXLIKE* out_dr_dq = nullptr,
		bool resize_out_dr_dq_to3x4 = true) const
	{
		using mrpt::square;
		using std::sqrt;

		if (out_dr_dq && resize_out_dr_dq_to3x4) out_dr_dq->setSize(3, 4);
		const T discr = r() * y() - x() * z();
		if (discr > 0.49999)
		{  // pitch = 90 deg
			pitch = 0.5 * M_PI;
			yaw = -2 * atan2(x(), r());
			roll = 0;
			if (out_dr_dq)
			{
				out_dr_dq->setZero();
				(*out_dr_dq)(0, 0) = +2 / x();
				(*out_dr_dq)(0, 2) = -2 * r() / (x() * x());
			}
		}
		else if (discr < -0.49999)
		{  // pitch =-90 deg
			pitch = -0.5 * M_PI;
			yaw = 2 * atan2(x(), r());
			roll = 0;
			if (out_dr_dq)
			{
				out_dr_dq->setZero();
				(*out_dr_dq)(0, 0) = -2 / x();
				(*out_dr_dq)(0, 2) = +2 * r() / (x() * x());
			}
		}
		else
		{  // Non-degenerate case:
			yaw = ::atan2(
				2 * (r() * z() + x() * y()), 1 - 2 * (y() * y() + z() * z()));
			pitch = ::asin(2 * discr);
			roll = ::atan2(
				2 * (r() * x() + y() * z()), 1 - 2 * (x() * x() + y() * y()));
			if (out_dr_dq)
			{
				// Auxiliary terms:
				const double val1 = (2 * x() * x() + 2 * y() * y() - 1);
				const double val12 = square(val1);
				const double val2 = (2 * r() * x() + 2 * y() * z());
				const double val22 = square(val2);
				const double xy2 = 2 * x() * y();
				const double rz2 = 2 * r() * z();
				const double ry2 = 2 * r() * y();
				const double val3 = (2 * y() * y() + 2 * z() * z() - 1);
				const double val4 =
					((square(rz2 + xy2) / square(val3) + 1) * val3);
				const double val5 = (4 * (rz2 + xy2)) / square(val3);
				const double val6 =
					1.0 / (square(rz2 + xy2) / square(val3) + 1);
				const double val7 = 2.0 / sqrt(1 - square(ry2 - 2 * x() * z()));
				const double val8 = (val22 / val12 + 1);
				const double val9 = -2.0 / val8;
				// row 1:
				(*out_dr_dq)(0, 0) = -2 * z() / val4;
				(*out_dr_dq)(0, 1) = -2 * y() / val4;
				(*out_dr_dq)(0, 2) = -(2 * x() / val3 - y() * val5) * val6;
				(*out_dr_dq)(0, 3) = -(2 * r() / val3 - z() * val5) * val6;
				// row 2:
				(*out_dr_dq)(1, 0) = y() * val7;
				(*out_dr_dq)(1, 1) = -z() * val7;
				(*out_dr_dq)(1, 2) = r() * val7;
				(*out_dr_dq)(1, 3) = -x() * val7;
				// row 3:
				(*out_dr_dq)(2, 0) = val9 * x() / val1;
				(*out_dr_dq)(2, 1) =
					val9 * (r() / val1 - (2 * x() * val2) / val12);
				(*out_dr_dq)(2, 2) =
					val9 * (z() / val1 - (2 * y() * val2) / val12);
				(*out_dr_dq)(2, 3) = val9 * y() / val1;
			}
		}
	}

	inline CQuaternion operator*(const T& factor)
	{
		CQuaternion q = *this;
		q *= factor;
		return q;
	}

};	// end class

/** A quaternion of data type "double" */
using CQuaternionDouble = CQuaternion<double>;
/** A quaternion of data type "float" */
using CQuaternionFloat = CQuaternion<float>;

}  // namespace mrpt::math
