/* +------------------------------------------------------------------------+
   |                     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 <mrpt/3rdparty/do_opencv_includes.h>
#include <mrpt/config.h>
#include <mrpt/vision/pnp_algos.h>
#include <mrpt/vision/utils.h>

#include <Eigen/Core>
#include <Eigen/Dense>
#include <iostream>
#if MRPT_HAS_OPENCV
#include <opencv2/core/eigen.hpp>
#endif

#include "dls.h"
#include "epnp.h"
#include "lhm.h"
#include "p3p.h"
#include "posit.h"
#include "ppnp.h"
#include "rpnp.h"
#include "upnp.h"

bool mrpt::vision::pnp::CPnP::dls(
    const Eigen::Ref<Eigen::MatrixXd> obj_pts,
    const Eigen::Ref<Eigen::MatrixXd> img_pts,
    int n,
    const Eigen::Ref<Eigen::MatrixXd> cam_intrinsic,
    Eigen::Ref<Eigen::MatrixXd> pose_mat)
{
  try
  {
#if MRPT_HAS_OPENCV == 1

    // Input 2d/3d correspondences and camera intrinsic matrix
    Eigen::MatrixXd cam_in_eig, img_pts_eig, obj_pts_eig;

    // Check for consistency of input matrix dimensions
    if (img_pts.rows() != obj_pts.rows() || img_pts.cols() != obj_pts.cols())
      throw(2);
    else if (cam_intrinsic.rows() != 3 || cam_intrinsic.cols() != 3)
      throw(3);

    if (obj_pts.rows() < obj_pts.cols())
    {
      cam_in_eig = cam_intrinsic.transpose();
      img_pts_eig = img_pts.transpose().block(0, 0, n, 2);
      obj_pts_eig = obj_pts.transpose();
    }
    else
    {
      cam_in_eig = cam_intrinsic;
      img_pts_eig = img_pts.block(0, 0, n, 2);
      obj_pts_eig = obj_pts;
    }

    // Output pose
    Eigen::Matrix3d R_eig;
    Eigen::MatrixXd t_eig;

    // Compute pose
    cv::Mat cam_in_cv(3, 3, CV_32F), img_pts_cv(2, n, CV_32F), obj_pts_cv(3, n, CV_32F),
        R_cv(3, 3, CV_32F), t_cv(3, 1, CV_32F);

    cv::eigen2cv(cam_in_eig, cam_in_cv);
    cv::eigen2cv(img_pts_eig, img_pts_cv);
    cv::eigen2cv(obj_pts_eig, obj_pts_cv);

    mrpt::vision::pnp::dls d(obj_pts_cv, img_pts_cv);
    bool ret = d.compute_pose(R_cv, t_cv);

    cv::cv2eigen(R_cv, R_eig);
    cv::cv2eigen(t_cv, t_eig);

    Eigen::Quaterniond q(R_eig);

    pose_mat << t_eig, q.vec();

    return ret;

#else
    throw(-1);
#endif
  }
  catch (int e)
  {
    switch (e)
    {
      case -1:
        std::cout << "Please install OpenCV for DLS-PnP" << std::endl;
        break;
      case 2:
        std::cout << "2d/3d correspondences mismatch\n Check dimension "
                     "of obj_pts and img_pts"
                  << std::endl;
        break;
      case 3:
        std::cout << "Camera intrinsic matrix does not have 3x3 dimensions " << std::endl;
        break;
    }
    return false;
  }
}

bool mrpt::vision::pnp::CPnP::epnp(
    const Eigen::Ref<Eigen::MatrixXd> obj_pts,
    const Eigen::Ref<Eigen::MatrixXd> img_pts,
    int n,
    const Eigen::Ref<Eigen::MatrixXd> cam_intrinsic,
    Eigen::Ref<Eigen::MatrixXd> pose_mat)
{
  try
  {
#if MRPT_HAS_OPENCV == 1

    // Input 2d/3d correspondences and camera intrinsic matrix
    Eigen::MatrixXd cam_in_eig, img_pts_eig, obj_pts_eig;

    // Check for consistency of input matrix dimensions
    if (img_pts.rows() != obj_pts.rows() || img_pts.cols() != obj_pts.cols())
      throw(2);
    else if (cam_intrinsic.rows() != 3 || cam_intrinsic.cols() != 3)
      throw(3);

    if (obj_pts.rows() < obj_pts.cols())
    {
      cam_in_eig = cam_intrinsic.transpose();
      img_pts_eig = img_pts.transpose().block(0, 0, n, 2);
      obj_pts_eig = obj_pts.transpose();
    }
    else
    {
      cam_in_eig = cam_intrinsic;
      img_pts_eig = img_pts.block(0, 0, n, 2);
      obj_pts_eig = obj_pts;
    }

    // Output pose
    Eigen::Matrix3d R_eig;
    Eigen::MatrixXd t_eig;

    // Compute pose
    cv::Mat cam_in_cv(3, 3, CV_32F), img_pts_cv(2, n, CV_32F), obj_pts_cv(3, n, CV_32F), R_cv, t_cv;

    cv::eigen2cv(cam_in_eig, cam_in_cv);
    cv::eigen2cv(img_pts_eig, img_pts_cv);
    cv::eigen2cv(obj_pts_eig, obj_pts_cv);

    mrpt::vision::pnp::epnp e(cam_in_cv, obj_pts_cv, img_pts_cv);
    e.compute_pose(R_cv, t_cv);

    cv::cv2eigen(R_cv, R_eig);
    cv::cv2eigen(t_cv, t_eig);

    Eigen::Quaterniond q(R_eig);

    pose_mat << t_eig, q.vec();

    return true;

#else
    throw(-1);
#endif
  }
  catch (int e)
  {
    switch (e)
    {
      case -1:
        std::cout << "Please install OpenCV for DLS-PnP" << std::endl;
        break;
      case 2:
        std::cout << "2d/3d correspondences mismatch\n Check dimension "
                     "of obj_pts and img_pts"
                  << std::endl;
        break;
      case 3:
        std::cout << "Camera intrinsic matrix does not have 3x3 dimensions " << std::endl;
        break;
    }
    return false;
  }
}

bool mrpt::vision::pnp::CPnP::upnp(
    const Eigen::Ref<Eigen::MatrixXd> obj_pts,
    const Eigen::Ref<Eigen::MatrixXd> img_pts,
    int n,
    const Eigen::Ref<Eigen::MatrixXd> cam_intrinsic,
    Eigen::Ref<Eigen::MatrixXd> pose_mat)
{
  try
  {
#if MRPT_HAS_OPENCV == 1

    // Input 2d/3d correspondences and camera intrinsic matrix
    Eigen::MatrixXd cam_in_eig, img_pts_eig, obj_pts_eig;

    // Check for consistency of input matrix dimensions
    if (img_pts.rows() != obj_pts.rows() || img_pts.cols() != obj_pts.cols())
      throw(2);
    else if (cam_intrinsic.rows() != 3 || cam_intrinsic.cols() != 3)
      throw(3);

    if (obj_pts.rows() < obj_pts.cols())
    {
      cam_in_eig = cam_intrinsic.transpose();
      img_pts_eig = img_pts.transpose().block(0, 0, n, 2);
      obj_pts_eig = obj_pts.transpose();
    }
    else
    {
      cam_in_eig = cam_intrinsic;
      img_pts_eig = img_pts.block(0, 0, n, 2);
      obj_pts_eig = obj_pts;
    }

    // Output pose
    Eigen::Matrix3d R_eig;
    Eigen::MatrixXd t_eig;

    // Compute pose
    cv::Mat cam_in_cv(3, 3, CV_32F), img_pts_cv(2, n, CV_32F), obj_pts_cv(3, n, CV_32F), R_cv, t_cv;

    cv::eigen2cv(cam_in_eig, cam_in_cv);
    cv::eigen2cv(img_pts_eig, img_pts_cv);
    cv::eigen2cv(obj_pts_eig, obj_pts_cv);

    mrpt::vision::pnp::upnp u(cam_in_cv, obj_pts_cv, img_pts_cv);
    u.compute_pose(R_cv, t_cv);

    cv::cv2eigen(R_cv, R_eig);
    cv::cv2eigen(t_cv, t_eig);

    Eigen::Quaterniond q(R_eig);

    pose_mat << t_eig, q.vec();

    return true;
#else
    throw(-1);
#endif
  }
  catch (int e)
  {
    switch (e)
    {
      case -1:
        std::cout << "Please install OpenCV for DLS-PnP" << std::endl;
        break;
      case 2:
        std::cout << "2d/3d correspondences mismatch\n Check dimension "
                     "of obj_pts and img_pts"
                  << std::endl;
        break;
      case 3:
        std::cout << "Camera intrinsic matrix does not have 3x3 dimensions " << std::endl;
        break;
    }
    return false;
  }
}

bool mrpt::vision::pnp::CPnP::p3p(
    const Eigen::Ref<Eigen::MatrixXd> obj_pts,
    const Eigen::Ref<Eigen::MatrixXd> img_pts,
    int n,
    const Eigen::Ref<Eigen::MatrixXd> cam_intrinsic,
    Eigen::Ref<Eigen::MatrixXd> pose_mat)
{
  try
  {
    // Input 2d/3d correspondences and camera intrinsic matrix
    Eigen::MatrixXd cam_in_eig, img_pts_eig, obj_pts_eig;

    // Check for consistency of input matrix dimensions
    if (img_pts.rows() != obj_pts.rows() || img_pts.cols() != obj_pts.cols())
      throw(2);
    else if (cam_intrinsic.rows() != 3 || cam_intrinsic.cols() != 3)
      throw(3);

    if (obj_pts.rows() < obj_pts.cols())
    {
      cam_in_eig = cam_intrinsic.transpose();
      img_pts_eig = img_pts.transpose().block(0, 0, n, 2);
      obj_pts_eig = obj_pts.transpose();
    }
    else
    {
      cam_in_eig = cam_intrinsic;
      img_pts_eig = img_pts.block(0, 0, n, 2);
      obj_pts_eig = obj_pts;
    }

    // Output pose
    Eigen::Matrix3d R;
    Eigen::Vector3d t;

    // Compute pose
    mrpt::vision::pnp::p3p p(cam_in_eig);
    bool ret = p.solve(R, t, obj_pts_eig, img_pts_eig);

    Eigen::Quaterniond q(R);

    pose_mat << t, q.vec();

    return ret;
  }
  catch (int e)
  {
    switch (e)
    {
      case 2:
        std::cout << "2d/3d correspondences mismatch\n Check dimension "
                     "of obj_pts and img_pts"
                  << std::endl;
        break;
      case 3:
        std::cout << "Camera intrinsic matrix does not have 3x3 dimensions " << std::endl;
        break;
    }
    return false;
  }
}

bool mrpt::vision::pnp::CPnP::rpnp(
    const Eigen::Ref<Eigen::MatrixXd> obj_pts,
    const Eigen::Ref<Eigen::MatrixXd> img_pts,
    int n,
    const Eigen::Ref<Eigen::MatrixXd> cam_intrinsic,
    Eigen::Ref<Eigen::MatrixXd> pose_mat)
{
  try
  {
    // Input 2d/3d correspondences and camera intrinsic matrix
    Eigen::MatrixXd cam_in_eig, img_pts_eig, obj_pts_eig;

    // Check for consistency of input matrix dimensions
    if (img_pts.rows() != obj_pts.rows() || img_pts.cols() != obj_pts.cols())
      throw(2);
    else if (cam_intrinsic.rows() != 3 || cam_intrinsic.cols() != 3)
      throw(3);

    if (obj_pts.rows() < obj_pts.cols())
    {
      cam_in_eig = cam_intrinsic.transpose();
      img_pts_eig = img_pts.transpose();
      obj_pts_eig = obj_pts.transpose();
    }
    else
    {
      cam_in_eig = cam_intrinsic;
      img_pts_eig = img_pts;
      obj_pts_eig = obj_pts;
    }

    // Output pose
    Eigen::Matrix3d R;
    Eigen::Vector3d t;

    // Compute pose
    mrpt::vision::pnp::rpnp r(obj_pts_eig, img_pts_eig, cam_in_eig, n);
    bool ret = r.compute_pose(R, t);

    Eigen::Quaterniond q(R);

    pose_mat << t, q.vec();

    return ret;
  }
  catch (int e)
  {
    switch (e)
    {
      case 2:
        std::cout << "2d/3d correspondences mismatch\n Check dimension "
                     "of obj_pts and img_pts"
                  << std::endl;
        break;
      case 3:
        std::cout << "Camera intrinsic matrix does not have 3x3 dimensions " << std::endl;
        break;
    }
    return false;
  }
}

bool mrpt::vision::pnp::CPnP::ppnp(
    const Eigen::Ref<Eigen::MatrixXd> obj_pts,
    const Eigen::Ref<Eigen::MatrixXd> img_pts,
    int n,
    const Eigen::Ref<Eigen::MatrixXd> cam_intrinsic,
    Eigen::Ref<Eigen::MatrixXd> pose_mat)
{
  try
  {
    // Input 2d/3d correspondences and camera intrinsic matrix
    Eigen::MatrixXd cam_in_eig, img_pts_eig, obj_pts_eig;

    // Check for consistency of input matrix dimensions
    if (img_pts.rows() != obj_pts.rows() || img_pts.cols() != obj_pts.cols())
      throw(2);
    else if (cam_intrinsic.rows() != 3 || cam_intrinsic.cols() != 3)
      throw(3);

    if (obj_pts.rows() < obj_pts.cols())
    {
      cam_in_eig = cam_intrinsic.transpose();
      img_pts_eig = img_pts.transpose();
      obj_pts_eig = obj_pts.transpose();
    }
    else
    {
      cam_in_eig = cam_intrinsic;
      img_pts_eig = img_pts;
      obj_pts_eig = obj_pts;
    }

    // Output pose
    Eigen::Matrix3d R;
    Eigen::Vector3d t;

    // Compute pose
    mrpt::vision::pnp::ppnp p(obj_pts_eig, img_pts_eig, cam_in_eig);

    bool ret = p.compute_pose(R, t, n);

    Eigen::Quaterniond q(R);

    pose_mat << t, q.vec();

    return ret;
  }
  catch (int e)
  {
    switch (e)
    {
      case 2:
        std::cout << "2d/3d correspondences mismatch\n Check dimension "
                     "of obj_pts and img_pts"
                  << std::endl;
        break;
      case 3:
        std::cout << "Camera intrinsic matrix does not have 3x3 dimensions " << std::endl;
        break;
    }
    return false;
  }
}

bool mrpt::vision::pnp::CPnP::posit(
    const Eigen::Ref<Eigen::MatrixXd> obj_pts,
    const Eigen::Ref<Eigen::MatrixXd> img_pts,
    int n,
    const Eigen::Ref<Eigen::MatrixXd> cam_intrinsic,
    Eigen::Ref<Eigen::MatrixXd> pose_mat)
{
  try
  {
    // Input 2d/3d correspondences and camera intrinsic matrix
    Eigen::MatrixXd cam_in_eig, img_pts_eig, obj_pts_eig;

    // Check for consistency of input matrix dimensions
    if (img_pts.rows() != obj_pts.rows() || img_pts.cols() != obj_pts.cols())
      throw(2);
    else if (cam_intrinsic.rows() != 3 || cam_intrinsic.cols() != 3)
      throw(3);

    if (obj_pts.rows() < obj_pts.cols())
    {
      cam_in_eig = cam_intrinsic.transpose();
      img_pts_eig = img_pts.transpose().block(0, 0, n, 2);
      obj_pts_eig = obj_pts.transpose();
    }
    else
    {
      cam_in_eig = cam_intrinsic;
      img_pts_eig = img_pts.block(0, 0, n, 2);
      obj_pts_eig = obj_pts;
    }

    // Output pose
    Eigen::Matrix3d R;
    Eigen::Vector3d t;

    // Compute pose
    mrpt::vision::pnp::posit p(obj_pts_eig, img_pts_eig, cam_in_eig, n);

    bool ret = p.compute_pose(R, t);

    Eigen::Quaterniond q(R);

    pose_mat << t, q.vec();

    return ret;
  }
  catch (int e)
  {
    switch (e)
    {
      case 2:
        std::cout << "2d/3d correspondences mismatch\n Check dimension "
                     "of obj_pts and img_pts"
                  << std::endl;
        break;
      case 3:
        std::cout << "Camera intrinsic matrix does not have 3x3 dimensions " << std::endl;
        break;
    }
    return false;
  }
}

bool mrpt::vision::pnp::CPnP::lhm(
    const Eigen::Ref<Eigen::MatrixXd> obj_pts,
    const Eigen::Ref<Eigen::MatrixXd> img_pts,
    int n,
    const Eigen::Ref<Eigen::MatrixXd> cam_intrinsic,
    Eigen::Ref<Eigen::MatrixXd> pose_mat)
{
  try
  {
    // Input 2d/3d correspondences and camera intrinsic matrix
    Eigen::MatrixXd cam_in_eig, img_pts_eig, obj_pts_eig;

    // Check for consistency of input matrix dimensions
    if (img_pts.rows() != obj_pts.rows() || img_pts.cols() != obj_pts.cols())
      throw(2);
    else if (cam_intrinsic.rows() != 3 || cam_intrinsic.cols() != 3)
      throw(3);

    if (obj_pts.rows() < obj_pts.cols())
    {
      cam_in_eig = cam_intrinsic.transpose();
      img_pts_eig = img_pts.transpose();
      obj_pts_eig = obj_pts.transpose();
    }
    else
    {
      cam_in_eig = cam_intrinsic;
      img_pts_eig = img_pts;
      obj_pts_eig = obj_pts;
    }

    // Output pose
    Eigen::Matrix3d R;
    Eigen::Vector3d t;

    // Compute pose
    mrpt::vision::pnp::lhm l(obj_pts_eig, img_pts_eig, cam_intrinsic, n);

    bool ret = l.compute_pose(R, t);

    Eigen::Quaterniond q(R);

    pose_mat << t, q.vec();

    return ret;
  }
  catch (int e)
  {
    switch (e)
    {
      case 2:
        std::cout << "2d/3d correspondences mismatch\n Check dimension "
                     "of obj_pts and img_pts"
                  << std::endl;
        break;
      case 3:
        std::cout << "Camera intrinsic matrix does not have 3x3 dimensions " << std::endl;
        break;
    }
    return false;
  }
}
