/* +------------------------------------------------------------------------+
   |                     Mobile Robot Programming Toolkit (MRPT)            |
   |                          https://www.mrpt.org/                         |
   |                                                                        |
   | Copyright (c) 2005-2024, 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 <iostream>

//#include <mrpt/math/types_math.h>  // Eigen must be included first via MRPT to
// enable the plugin system
#include <Eigen/Dense>
#include <Eigen/SVD>
#include <Eigen/StdVector>

#include "lhm.h"
using namespace mrpt::vision::pnp;

lhm::lhm(Eigen::MatrixXd obj_pts_, Eigen::MatrixXd img_pts_, Eigen::MatrixXd cam_, int n0) : F(n0)
{
  obj_pts = obj_pts_;
  img_pts = img_pts_;
  cam_intrinsic = cam_;
  n = n0;

  // Store obj_pts as 3XN and img_projections as 2XN matrices
  P = obj_pts.transpose();
  Q = Eigen::MatrixXd::Ones(3, n);

  Q = img_pts.transpose();

  t.setZero();
}

void lhm::estimate_t()
{
  Eigen::Vector3d sum_;
  sum_.setZero();
  for (int i = 0; i < n; i++) sum_ += F[i] * R * P.col(i);
  t = G * sum_;
}

void lhm::xform()
{
  for (int i = 0; i < n; i++) Q.col(i) = R * P.col(i) + t;
}

Eigen::Matrix4d lhm::qMatQ(Eigen::VectorXd q)
{
  Eigen::Matrix4d Q_(4, 4);

  Q_ << q(0), -q(1), -q(2), -q(3), q(1), q(0), -q(3), q(2), q(2), q(3), q(0), -q(1), q(3), -q(2),
      q(1), q(0);

  return Q_;
}

Eigen::Matrix4d lhm::qMatW(Eigen::VectorXd q)
{
  Eigen::Matrix4d Q_(4, 4);

  Q_ << q(0), -q(1), -q(2), -q(3), q(1), q(0), q(3), -q(2), q(2), -q(3), q(0), q(1), q(3), q(2),
      -q(1), q(0);

  return Q_;
}

void lhm::absKernel()
{
  for (int i = 0; i < n; i++) Q.col(i) = F[i] * Q.col(i);

  Eigen::Vector3d P_bar, Q_bar;
  P_bar = P.rowwise().mean();
  Q_bar = Q.rowwise().mean();

  for (int i = 0; i < n; i++)
  {
    P.col(i) = P.col(i) - P_bar;
    Q.col(i) = Q.col(i) - Q_bar;
  }

  //<------------------- Use SVD Solution ------------------->//
  /*
  Eigen::Matrix3d M;
  M.setZero();

  for (i = 0; i < n; i++)
    M += P.col(i)*Q.col(i).transpose();

  Eigen::JacobiSVD<Eigen::MatrixXd> svd(M, Eigen::ComputeThinU |
  Eigen::ComputeThinV);

  R = svd.matrixV()*svd.matrixU().transpose();

  Eigen::Matrix3d dummy;
  dummy.col(0) = -svd.matrixV().col(0);
  dummy.col(1) = -svd.matrixV().col(1);
  dummy.col(2) = svd.matrixV().col(2);

  if (R.determinant() == 1)
  {
    estimate_t();
    if (t(2) < 0)
    {
      R = dummy*svd.matrixU().transpose();
      estimate_t();
    }
  }
  else
  {
    R = -dummy*svd.matrixU().transpose();
    estimate_t();
    if (t(2) < 0)
    {
      R = -svd.matrixV()*svd.matrixU().transpose();
      estimate_t();
    }
  }

  err2 = 0;
  xform();

  Eigen::Vector3d vec;
  Eigen::Matrix3d I3 = Eigen::MatrixXd::Identity(3, 3);

  for (i = 0; i < n; i++)
  {
    vec = (I3 - F[i])*Q.col(i);
    err2 += vec.squaredNorm();
  }
  */
  //<------------------- Use QTN Solution ------------------->//

  Eigen::Matrix4d A;
  A.setZero();

  for (int i = 0; i < n; i++)
  {
    Eigen::Vector4d q1, q2;
    q1 << 1, Q.col(i);
    q2 << 1, P.col(i);
    A += qMatQ(q1).transpose() * qMatW(q2);
  }

  Eigen::EigenSolver<Eigen::Matrix4d> es(A);

  const Eigen::Matrix4d Ae = es.pseudoEigenvalueMatrix();
  Eigen::Vector4d D;  // Ae.diagonal(); for some reason this leads to an
  // internal compiler error in MSVC11... (sigh)
  for (int i = 0; i < 4; i++) D[i] = Ae(i, i);

  Eigen::Matrix4d V_mat = es.pseudoEigenvectors();

  Eigen::Vector4d::Index max_index;

  D.maxCoeff(&max_index);

  Eigen::Vector4d V;

  V = V_mat.col(max_index);

  Eigen::Quaterniond q(V(0), V(1), V(2), V(3));

  R = q.toRotationMatrix();

  estimate_t();

  err2 = 0;
  xform();

  Eigen::Vector3d vec;
  Eigen::Matrix3d I3 = Eigen::MatrixXd::Identity(3, 3);

  for (int i = 0; i < n; i++)
  {
    vec = (I3 - F[i]) * Q.col(i);
    err2 += vec.squaredNorm();
  }
}

bool lhm::compute_pose(Eigen::Ref<Eigen::Matrix3d> R_, Eigen::Ref<Eigen::Vector3d> t_)
{
  int i, j = 0;

  Eigen::VectorXd p_bar;
  Eigen::Matrix3d sum_F, I3;
  I3 = Eigen::MatrixXd::Identity(3, 3);
  sum_F.setZero();

  p_bar = P.rowwise().mean();

  for (i = 0; i < n; i++)
  {
    P.col(i) -= p_bar;
    F[i] = Q.col(i) * Q.col(i).transpose() / Q.col(i).squaredNorm();
    sum_F = sum_F + F[i];
  }

  G = (I3 - sum_F / n).inverse() / n;

  err = 0;
  err2 = 1000;
  absKernel();

  while (std::abs(err2 - err) > TOL_LHM && err2 > EPSILON_LHM)
  {
    err = err2;

    absKernel();

    j += 1;
    if (j > 100) break;
  }

  R_ = R;
  t_ = t - R * p_bar;

  return true;
}
