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

#include "vision-precomp.h"	 // Precompiled headers
//
#include <Eigen/Dense>
#include <Eigen/SVD>
#include <iostream>

#include "posit.h"

mrpt::vision::pnp::posit::posit(
	Eigen::MatrixXd obj_pts_, Eigen::MatrixXd img_pts_,
	Eigen::MatrixXd camera_intrinsic_, int n0)
{
	obj_pts = obj_pts_;
	img_pts = img_pts_.block(0, 0, n0, 2);
	cam_intrinsic = camera_intrinsic_;
	R = Eigen::MatrixXd::Identity(3, 3);
	t = Eigen::VectorXd::Zero(3);
	f = (cam_intrinsic(0, 0) + cam_intrinsic(1, 1)) / 2;

	obj_matrix =
		(obj_pts.transpose() * obj_pts).inverse() * obj_pts.transpose();

	n = n0;

	obj_vecs = Eigen::MatrixXd::Zero(n0, 3);

	for (int i = 0; i < n; i++)
		obj_vecs.row(i) = obj_pts.row(i) - obj_pts.row(0);

	img_vecs = Eigen::MatrixXd::Zero(n0, 2);
	img_vecs_old = img_vecs;

	epsilons = Eigen::VectorXd::Zero(n);
}

void mrpt::vision::pnp::posit::POS()
{
	Eigen::Vector3d I0, J0, r1, r2, r3;
	double I0_norm, J0_norm;

	int i;
	double scale;

	for (i = 0; i < 3; i++)
	{
		I0(i) = obj_matrix.row(i).dot(img_vecs.col(0));
		J0(i) = obj_matrix.row(i).dot(img_vecs.col(1));
	}

	I0_norm = I0.norm();
	J0_norm = J0.norm();

	scale = (I0_norm + J0_norm) / 2;

	/*Computing TRANSLATION */
	t(0) = img_pts(0, 0) / scale;
	t(1) = img_pts(0, 1) / scale;
	t(2) = f / scale;

	/* Computing ROTATION */
	r1 = I0 / I0_norm;
	r2 = J0 / J0_norm;
	r3 = r1.cross(r2);

	R.row(0) = r1;
	R.row(1) = r2;
	R.row(2) = r3;
}

/**
Iterate over results obtained by the POS function;
see paper "Model-Based Object Pose in 25 Lines of Code", IJCV 15, pp. 123-141,
1995.
*/
bool mrpt::vision::pnp::posit::compute_pose(
	Eigen::Ref<Eigen::Matrix3d> R_, Eigen::Ref<Eigen::Vector3d> t_)
{
	Eigen::FullPivLU<Eigen::MatrixXd> lu(obj_pts);
	if (lu.rank() < 3) return false;

	int i, iCount;
	long imageDiff = 1000;

	for (iCount = 0; iCount < LOOP_MAX_COUNT; iCount++)
	{
		if (iCount == 0)
		{
			for (i = 0; i < img_vecs.rows(); i++)
				img_vecs.row(i) = img_pts.row(i) - img_pts.row(0);
		}

		else
		{
			// Compute new image vectors
			epsilons.setZero();
			for (i = 0; i < n; i++)
			{
				epsilons(i) += obj_vecs.row(i).dot(R.row(2));
			}
			epsilons /= t(2);

			// Corrected image vectors
			for (i = 0; i < n; i++)
			{
				img_vecs.row(i) =
					img_pts.row(i) * (1 + epsilons(i)) - img_pts.row(0);
			}

			imageDiff = this->get_img_diff();
		}

		img_vecs_old = img_vecs;

		this->POS();

		if (iCount > 0 && imageDiff == 0) break;

		if (iCount == LOOP_MAX_COUNT)
		{
			std::cout << "Solution Not converged" << std::endl << std::endl;
			break;
		}
	}
	R_ = R;
	t_ = t;

	return true;
}

long mrpt::vision::pnp::posit::get_img_diff()
{
	double sumOfDiffs = 0;

	for (int i = 0; i < n; i++)
	{
		for (int j = 0; j < 2; j++)
		{
			sumOfDiffs += std::abs(
				floor(0.5 + img_vecs(i, j)) - floor(0.5 + img_vecs_old(i, j)));
		}
	}
	return static_cast<long>(sumOfDiffs);
}
