#include <Eigen/Dense>
#include <cstdint>//on CUDA, use native primitives instead (same for sqrt, abs, max and other std:: functions)
#define hostdevice//on CUDA, __host__ __device__

typedef int32_t index_t;

template<typename scalar_t> hostdevice inline scalar_t P1PGradSmallValue() { return scalar_t(1e-10); }
template<> hostdevice inline float P1PGradSmallValue<float>() { return 1e-6f; }

namespace P1AG
{
	//helper function, solves the canonical P1AC problem (one correspondence, two gradients) when the depth of the point is known (mainly, computes the up to two solutions with the sign of the epsilons)
	template <typename point2D_t, typename grad2D_t, typename grad3D_t, typename simplemat4_t, typename mat4_t>
	hostdevice void CanonicalP1AC_solveknownd(index_t nb_keep, index_t& current_stored, typename point2D_t::Scalar gu1sqnorm, typename point2D_t::Scalar gu2sqnorm, typename point2D_t::Scalar gx1sqnorm, typename point2D_t::Scalar gx2sqnorm, typename point2D_t::Scalar d,
		const point2D_t& p2d, const grad2D_t& g2d_1, const grad2D_t& g2d_2, const grad3D_t& g3d_1, const grad3D_t& g3d_2, const simplemat4_t& virtualcam_to_old,
		mat4_t* m, typename point2D_t::Scalar* store_d)
	{
		using scalar_t = typename point2D_t::Scalar;

		//build the matrix of modified 2D gradients
		Eigen::Matrix<scalar_t, 3, 3> migrad3;
		migrad3.template block<1, 3>(0, 0) = g3d_1.transpose();
		migrad3.template block<1, 3>(1, 0) = g3d_2.transpose();
		migrad3.template block<1, 3>(2, 0) = migrad3.template block<1, 3>(0, 0).cross(migrad3.template block<1, 3>(1, 0));//due to the way we write these matrices, we force R to be a rotation (det > 0)
		migrad3 = migrad3.inverse().eval();

		//build the matrix of modified canonical gradients with epsilon as third component
		Eigen::Matrix<scalar_t, 3, 3> mgrad2;
		mgrad2.template block<1, 2>(0, 0) = g2d_1.transpose() * d;
		mgrad2.template block<1, 2>(1, 0) = g2d_2.transpose() * d;
		mgrad2(0, 2) = gu1sqnorm - d * d * gx1sqnorm;//epsilon1²
		bool has_epsilon1 = true;//is epsilon1² > 0?
		bool has_epsilon2 = true;//is epsilon2² > 0?
		if (mgrad2(0, 2) <= 0.0f)//epsilon1² <= 0.0f --> impossible, but OK with numerical errors
		{
			if (-mgrad2(0, 2) < P1PGradSmallValue<scalar_t>() * gx1sqnorm)//consider epsilon1 = 0
			{
				mgrad2(0, 2) = 0.0f;
				has_epsilon1 = false;//the only choice is epsilon1 = 0
			}
			else
				return;//invalid solution
		}
		else
			mgrad2(0, 2) = std::sqrt(mgrad2(0, 2));
		mgrad2(1, 2) = gu2sqnorm - d * d * gx2sqnorm;//epsilon2²
		if (mgrad2(1, 2) <= 0.0f)//epsilon2² <= 0.0f --> impossible, but OK with numerical errors
		{
			if (-mgrad2(1, 2) < P1PGradSmallValue<scalar_t>() * gx2sqnorm)
			{
				mgrad2(1, 2) = 0.0f;
				has_epsilon2 = false;//the only choice is epsilon2 = 0
			}
			else
				return;//invalid solution
		}
		else
		{
			mgrad2(1, 2) = std::sqrt(mgrad2(1, 2));
			if ((g3d_1.dot(g3d_2) - g2d_1.dot(g2d_2) * d * d > 0.0f) != (mgrad2(0, 2) * mgrad2(1, 2) > 0.0f))
			{
				mgrad2(1, 2) = -mgrad2(1, 2);//the sign of eps_1 * eps_2 is wrong, let's invert eps_2
			}
		}

		Eigen::Vector<scalar_t, 3> scaled_p2d;
		scaled_p2d.template block<2, 1>(0, 0) = p2d * d;
		scaled_p2d(2, 0) = d;

		mgrad2.template block<1, 3>(2, 0) = mgrad2.template block<1, 3>(0, 0).cross(mgrad2.template block<1, 3>(1, 0));
		Eigen::Matrix<scalar_t, 4, 4> pose_hyp = Eigen::Matrix<scalar_t, 4, 4>::Identity();
		pose_hyp.template block<3, 3>(0, 0) = migrad3 * mgrad2;//++
		//if (CheckMatrixIsRotation<scalar_t>(pose_hyp))//it is a rotation unless our math is incorrect
		//{
		pose_hyp.template block<3, 1>(0, 3) = scaled_p2d - pose_hyp.template block<3, 3>(0, 0) * Eigen::Vector<scalar_t, 3>(scalar_t(0), scalar_t(0), scalar_t(1));
		m[current_stored] = pose_hyp * virtualcam_to_old;//new_to_virtualcam * virtualcam_to_old = new_to_old
		store_d[current_stored] = d;
		if (++current_stored == nb_keep)//check in case the user only wants one output
			return;
		//}

		if (has_epsilon1 || has_epsilon2)
		{
			//tweak the matrix of canonical gradients to get another solution
			mgrad2(0, 2) = -mgrad2(0, 2);//get the other root for epsilon_1
			mgrad2(1, 2) = -mgrad2(1, 2);//but also for epsilon_2
			mgrad2.template block<1, 3>(2, 0) = mgrad2.template block<1, 3>(0, 0).cross(mgrad2.template block<1, 3>(1, 0));
			pose_hyp = Eigen::Matrix<scalar_t, 4, 4>::Identity();
			pose_hyp.template block<3, 3>(0, 0) = migrad3 * mgrad2;//-+
			//if (CheckMatrixIsRotation<scalar_t>(pose_hyp))//it is a rotation unless our math is incorrect
			//{
			pose_hyp.template block<3, 1>(0, 3) = scaled_p2d - pose_hyp.template block<3, 3>(0, 0) * Eigen::Vector<scalar_t, 3>(scalar_t(0), scalar_t(0), scalar_t(1));
			m[current_stored] = pose_hyp * virtualcam_to_old;//new_to_virtualcam * virtualcam_to_old = new_to_old
			store_d[current_stored] = d;
			++current_stored;
			//if (++current_stored == nb_keep)
			//	return;
		//}
		}
	}

	/* main canonical P1AC solver
	\param[in] p2d image point U_a = (u_a, v_a)
	\param[in] g2d_1 first image gradient (2 components), something like Eigen::Vector2f/2d or Eigen::Map<Eigen::Vector2f/2d>
	\param[in] g2d_2 second image gradient (2 components), something like Eigen::Vector2f/2d or Eigen::Map<Eigen::Vector2f/2d>
	\param[in] gx1 first reference gradient in canonical frame (2 components), something like Eigen::Vector2f/2d or Eigen::Map<Eigen::Vector2f/2d>
	\param[in] gx2 first reference gradient in canonical frame (2 components), something like Eigen::Vector2f/2d or Eigen::Map<Eigen::Vector2f/2d>
	\param[in] virtualcam_to_old de-canonicalization matrix (Eigen:::Matrix4f/4d)
	\param[out] m preallocated vector of matrices for output transformations (typically, size = 2), something like Eigen::Matrix4f/4d[2]
	\param[out] depths preallocated vector of pointdepths (size = size of m), something like float[2]
	\param[in] nb_keep size of m and depths, typically 2
	
	\return number of solutions (0 = unable to solve, 1 or 2) - only solutions with positive depth
	
	note: preallocation of output to be compatible with environments witout dynamic allocation (CUDA)
	*/
	template <typename point2D_t, typename grad2D_t, typename grad2Dx_t, typename simplemat4_t, typename mat4_t>
	hostdevice index_t CanonicalP1AC(const point2D_t& p2d, const grad2D_t& g2d_1, const grad2D_t& g2d_2, const grad2Dx_t& gx1, const grad2Dx_t& gx2,
		const simplemat4_t& virtualcam_to_old, mat4_t* m, typename point2D_t::Scalar* depths, index_t nb_keep = 2)
	{
		using scalar_t = typename point2D_t::Scalar;

		Eigen::Vector<scalar_t, 3> gu1(g2d_1[0], g2d_1[1], -g2d_1.dot(p2d));//G_1^T in the paper
		Eigen::Vector<scalar_t, 3> gu2(g2d_2[0], g2d_2[1], -g2d_2.dot(p2d));//G_2^T in the paper

		//(dgx^T epsilon) = gu^T.R (written as column vectors)

		const scalar_t gx1sqnorm = gx1.squaredNorm();
		const scalar_t gx2sqnorm = gx2.squaredNorm();
		const scalar_t gxdot = gx1.dot(gx2);
		const scalar_t gx1sqnormgx2sqnorm = gx1sqnorm * gx2sqnorm;//putting it here so on GPU, a lot of time is available before the result is actually needed
		const scalar_t gu1sqnorm = gu1.squaredNorm();
		const scalar_t gu2sqnorm = gu2.squaredNorm();
		const scalar_t gudot = gu1.dot(gu2);
		
		const scalar_t A = gx1sqnormgx2sqnorm - gxdot * gxdot;//A
		if (A <= P1PGradSmallValue<scalar_t>() * gx1sqnormgx2sqnorm)
			return 0;//3D gradients are colinear, meaning there is no solution. We refuse to solve in this case
		//using Ax^2 + 2Bx + C = 0
		const scalar_t Br = gxdot * gudot - scalar_t(0.5) * (gx1sqnorm * gu2sqnorm + gx2sqnorm * gu1sqnorm);//B'=B/2
		const scalar_t C = gu1sqnorm * gu2sqnorm - gudot * gudot;//C
		//if (C <= P1PGradSmallValue<scalar_t>() * gu1sqnorm)
		//	return 0;//2d augmented gradients are colinear, meaning no solution
		const scalar_t bsquare = Br * Br;
		const scalar_t ac = A * C;
		scalar_t discriminant = bsquare - ac;
		const scalar_t maxterm = (std::max)(std::abs(bsquare), std::abs(ac));
		index_t nbvalid = 0;
		scalar_t d = scalar_t(0.0);
		if (discriminant < scalar_t(0.0))//discriminant is probably -1e-10 or something like that
		{
			if (-discriminant < P1PGradSmallValue<scalar_t>() * maxterm)//epsilon error, tolerated
			{
				d = -Br / A;//set discriminant to 0
				CanonicalP1AC_solveknownd(nb_keep, nbvalid, gu1sqnorm, gu2sqnorm, gx1sqnorm, gx2sqnorm, d, p2d, gx1, gx2, gu1, gu2, virtualcam_to_old, m, depths);
			}
			//else: this is a serious "impossible" failure: return 0 solutions.
		}
		else//in a perfect world without noise, we would only need this branch
		{
			//scalar_t sol1 = (-Br - sqrt(discriminant)) / A;
			scalar_t sol2 = (-Br - std::sqrt(discriminant)) / A;//this root cannot be < 0
			if (sol2 <= scalar_t(0.0))//one solution for d, with both roots equal to zero. d = 0 does not allow solving. d < 0 is just some numerical noise
				return 0;
			else
			{
				d = std::sqrt(sol2);
				CanonicalP1AC_solveknownd(nb_keep, nbvalid, gu1sqnorm, gu2sqnorm, gx1sqnorm, gx2sqnorm, d, p2d, gx1, gx2, gu1, gu2, virtualcam_to_old, m, depths);
			}
		}
		return nbvalid;
	}

	/* main P1PGrad solver
	\param[out] virtualcam_to_3Dmodel: de-canonicalization matrix (filled by the function)
	\param[in] p2d image point U_a = (u_a, v_a), something like Eigen::Vector2f/2d or Eigen::Map<const Eigen::Vector2f/2d>
	\param[in] p3d reference point (x,y,z) where U_b = (x/z, y/z), something like Eigen::Vector3f/3d or Eigen::Map<const Eigen::Vector3f/3d>
	\param[in] g2d_1 first image gradient (2 components), something like Eigen::Vector2f/2d or Eigen::Map<const Eigen::Vector2f/2d>
	\param[in] g2d_2 second image gradient (2 components), something like Eigen::Vector2f/2d or Eigen::Map<const Eigen::Vector2f/2d>
	\param[in] g3d_1 first reference gradient (3 components), something like Eigen::Vector3f/3d or Eigen::Map<const Eigen::Vector3f/3d>
	\param[in] g3d_1 first reference gradient (3 components), something like Eigen::Vector3f/3d or Eigen::Map<const Eigen::Vector3f/3d>
	\param[out] m preallocated vector of matrices for output transformations (typically, size = 2), something like Eigen::Matrix4f/4d[2]
	\param[out] depths preallocated vector of point depths (size = size of m), something like float[2]
	\param[in] nb_keep size of m and depths, typically 2
	
	\return number of solutions (0 = unable to solve, 1 or 2) - only solutions with positive depth
	
	note: preallocation of output to be compatible with environments witout dynamic allocation (CUDA)
	*/
	template <typename point2D_t, typename point3D_t, typename grad2D_t, typename grad3D_t, typename mat4v_t>
	hostdevice index_t P1PGrad(Eigen::Matrix<typename point2D_t::Scalar, 4, 4>& virtualcam_to_3Dmodel, const point2D_t& p2d, const point3D_t& p3d, const grad2D_t& g2d_1, const grad2D_t& g2d_2, const grad3D_t& g3d_1, const grad3D_t& g3d_2,
		mat4v_t* m, typename point2D_t::Scalar* depths, index_t nb_keep = 2)
	{
		using scalar_t = typename point2D_t::Scalar;

		virtualcam_to_3Dmodel = Eigen::Matrix<scalar_t, 4, 4>::Identity();
		//we have to invert the matrix, which we do by transposing the rotation and having t -> -R^T.t
		virtualcam_to_3Dmodel.template block<1, 3>(0, 0) = g3d_1.normalized().transpose();
		virtualcam_to_3Dmodel.template block<1, 3>(2, 0) = g3d_1.cross(g3d_2).normalized().transpose();
		virtualcam_to_3Dmodel.template block<1, 3>(1, 0) = virtualcam_to_3Dmodel.template block<1, 3>(2, 0).cross(virtualcam_to_3Dmodel.template block<1, 3>(0, 0));
		virtualcam_to_3Dmodel.template block<3, 1>(0, 3) = virtualcam_to_3Dmodel.template block<3, 3>(0, 0) * (virtualcam_to_3Dmodel.template block<1, 3>(2, 0).transpose() - p3d);
		//this matrix sends (0 0 1)^T on X and the z=1 plane on the local tangent plane defined by X and the two 2D gradients

		//we are removing a component from the 3D vector and adding one to the 2D vector. Such a good joke!
		Eigen::Vector<scalar_t, 2> gx1 = (virtualcam_to_3Dmodel.template block<3, 3>(0, 0) * g3d_1).template block<2, 1>(0, 0);//3D gradients in the local reference frame
		Eigen::Vector<scalar_t, 2> gx2 = (virtualcam_to_3Dmodel.template block<3, 3>(0, 0) * g3d_2).template block<2, 1>(0, 0);
		return CanonicalP1AC(p2d, g2d_1, g2d_2, gx1, gx2, virtualcam_to_3Dmodel, m, depths, nb_keep);
	}
	
	/* main P1AG solver (designed so that its interface matches P1AC's)
	\param[in] p2d image point U_a = (u_a, v_a), something like Eigen::Vector2f/2d or Eigen::Map<const Eigen::Vector2f/2d>
	\param[in] p2d_src reference point U_b = (u_b, v_b), something like Eigen::Vector2f/2d or Eigen::Map<const Eigen::Vector2f/2d>
	\param[in] d_src depth of the reference point U_b, so that the 3D point is (u_b.d_src, v_b.d_src, d_src)
	\param[in] n local surface normal at point (u_b.d_src, v_b.d_src, d_src), something like Eigen::Vector3f/3d or Eigen::Map<const const Eigen::Vector3f/3d>
	\param[in] A affine matrix
	\param[out] m preallocated vector of matrices for output transformations (typically, size = 2), something like Eigen::Matrix4f/4d[2]
	\param[out] depths preallocated vector of point depths (size = size of m), something like float[2]
	\param[in] nb_keep size of m and depths, typically 2
	
	\return number of solutions (0 = unable to solve, 1 or 2) - only solutions with positive depth
	
	note: preallocation of output to be compatible with environments witout dynamic allocation (CUDA)
	*/
	template<typename point2D_t, typename normal_t, typename mat2_t, typename mat4v_t>
	hostdevice index_t P1AG(const point2D_t& p2d, const point2D_t& p2d_src, const typename point2D_t::Scalar d_src, const normal_t& n, const mat2_t& A, mat4v_t* m, typename point2D_t::Scalar* depths, index_t nb_keep = 2)
	{
		using scalar_t = typename point2D_t::Scalar;
		Eigen::Matrix<scalar_t, 4, 4> virtualcam_to_old = Eigen::Matrix<scalar_t, 4, 4>::Identity();
		virtualcam_to_old.template block<1, 3>(2, 0) = n.transpose();
		if (std::abs(n[0]) > std::abs(n[1]))//n has a large x component, use ~y as orthovector
			virtualcam_to_old.template block<1, 3>(0, 0) = Eigen::RowVector<scalar_t, 3>(scalar_t(0), scalar_t(1), scalar_t(0)).cross(virtualcam_to_old.template block<1, 3>(2, 0)).normalized();
		else//n has a large y component, use ~x as orthovector
			virtualcam_to_old.template block<1, 3>(0, 0) = Eigen::RowVector<scalar_t, 3>(scalar_t(1), scalar_t(0), scalar_t(0)).cross(virtualcam_to_old.template block<1, 3>(2, 0)).normalized();
		virtualcam_to_old.template block<1, 3>(1, 0) = virtualcam_to_old.template block<1, 3>(2, 0).cross(virtualcam_to_old.template block<1, 3>(0, 0));
		virtualcam_to_old.template block<3, 1>(0, 3) = virtualcam_to_old.template block<3, 3>(0, 0) * (virtualcam_to_old.template block<1, 3>(2, 0).transpose() - Eigen::Vector3<scalar_t>(p2d_src[0] * d_src, p2d_src[1] * d_src, d_src));
		//this matrix transforms points from old to virtualcam
		//now, we have to find the gradients
		//we have partial n/partial v = partial n / partial o . partial o / partial v, with A = partial n / partial o
		Eigen::Matrix<scalar_t, 2, 2> newA;
		newA(0, 0) = virtualcam_to_old(0, 0) - virtualcam_to_old(0, 2) * p2d_src[0];//newA <- partial o / partial v
		newA(0, 1) = virtualcam_to_old(1, 0) - virtualcam_to_old(1, 2) * p2d_src[0];
		newA(1, 0) = virtualcam_to_old(0, 1) - virtualcam_to_old(0, 2) * p2d_src[1];
		newA(1, 1) = virtualcam_to_old(1, 1) - virtualcam_to_old(1, 2) * p2d_src[1];
		newA = 1 / d_src * A * newA;
		//now, we have : grad2d_virtual = grad2d_new * newA. Written with column vectors, grad2d_virtual^T = newA^T * grad2d_new^T
		const Eigen::Vector<scalar_t, 2> g2d_1(scalar_t(1), scalar_t(0));
		const Eigen::Vector<scalar_t, 2> g2d_2(scalar_t(0), scalar_t(1));
		const Eigen::Vector<scalar_t, 2> gx1 = newA.template block<1, 2>(0, 0).transpose();//newA.transpose() * g2d_1;
		const Eigen::Vector<scalar_t, 2> gx2 = newA.template block<1, 2>(1, 0).transpose();//newA.transpose() * g2d_2;
		return CanonicalP1AC(p2d, g2d_1, g2d_2, gx1, gx2, virtualcam_to_old, m, depths, nb_keep);
	}
}