/* +------------------------------------------------------------------------+
   |                     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 "tfest-precomp.h"  // Precompiled headers
//
#include <mrpt/core/round.h>
#include <mrpt/math/utils.h>  // linspace()
#include <mrpt/poses/CPose3D.h>
#include <mrpt/poses/CPose3DQuat.h>
#include <mrpt/poses/Lie/SO.h>
#include <mrpt/random.h>
#include <mrpt/tfest/se3.h>

#include <iostream>
#include <numeric>

using namespace mrpt;
using namespace mrpt::tfest;
using namespace mrpt::random;
using namespace mrpt::poses;
using namespace mrpt::math;
using namespace std;

/*---------------------------------------------------------------
             se3_l2_robust
  ---------------------------------------------------------------*/
bool tfest::se3_l2_robust(
    const mrpt::tfest::TMatchingPairList& in_correspondences,
    const TSE3RobustParams& params,
    TSE3RobustResult& results)
{
  MRPT_START

  const size_t nCorrs = in_correspondences.size();

  // -------------------------------------------
  // RANSAC parameters
  // -------------------------------------------
  // Minimum error achieved so far
  double min_err = std::numeric_limits<double>::max();
  size_t max_size = 0;  // Maximum size of the consensus set so far
  double scale;         // Output scale

  // Minimum number of points to fit the model
  const size_t n = params.ransac_minSetSize;
  // Minimum number of points to be considered a good set
  const size_t d = mrpt::round(nCorrs * params.ransac_maxSetSizePct);
  // Maximum number of iterations
  const size_t maxIters = params.ransac_nmaxSimulations;

  ASSERTMSG_(
      d >= n,
      "Minimum number of points to be considered a good set is < Minimum "
      "number of points to fit the model");

  // -------------------------------------------
  // MAIN loop
  // -------------------------------------------
  for (size_t iterations = 0; iterations < maxIters; iterations++)
  {
    if (params.verbose)
      std::cout << "[tfest::se3_l2_robust] Iteration " << (iterations + 1) << "/" << maxIters
                << "\n";

    // Generate maybe inliers
    const auto rub = mrpt::math::linspace<uint32_t>(0, nCorrs - 1, nCorrs);
    const auto mbSet = getRandomGenerator().permuteVector(rub);

    std::vector<uint32_t> cSet;  // consensus set

    // Compute first inliers output
    TMatchingPairList mbInliers;
    mbInliers.reserve(n);
    for (size_t i = 0; mbInliers.size() < n && i < nCorrs; i++)
    {
      const size_t idx = mbSet[i];

      // User-provided filter:
      if (params.user_individual_compat_callback)
      {
        mrpt::tfest::TPotentialMatch pm;
        pm.idx_this = in_correspondences[idx].globalIdx;
        pm.idx_other = in_correspondences[idx].localIdx;
        if (!params.user_individual_compat_callback(pm)) continue;  // Skip this one!
      }

      mbInliers.push_back(in_correspondences[idx]);
      cSet.push_back(idx);
    }

    // Check minimum number:
    if (cSet.size() < n)
    {
      if (params.verbose)
        std::cerr << "[tfest::se3_l2_robust] Iter " << iterations
                  << ": It was not possible to find the min no of "
                     "(compatible) matching pairs.\n";
      continue;  // Try again
    }

    CPose3DQuat mbOutQuat;
    bool res = mrpt::tfest::se3_l2(mbInliers, mbOutQuat, scale, params.forceScaleToUnity);
    if (!res)
    {
      std::cerr << "[tfest::se3_l2_robust] tfest::se3_l2() returned "
                   "false for tentative subset during RANSAC "
                   "iteration!\n";
      continue;
    }

    // Maybe inliers Output
    const CPose3D mbOut = CPose3D(mbOutQuat);
    CVectorDouble mbOut_vec(7);
    mbOut_vec[0] = mbOut.x();
    mbOut_vec[1] = mbOut.y();
    mbOut_vec[2] = mbOut.z();

    mbOut_vec[3] = mbOut.yaw();
    mbOut_vec[4] = mbOut.pitch();
    mbOut_vec[5] = mbOut.roll();

    mbOut_vec[6] = scale;

    // Inner loop: for each point NOT in the maybe inliers
    for (size_t k = n; k < nCorrs; k++)
    {
      const size_t idx = mbSet[k];

      // User-provided filter:
      if (params.user_individual_compat_callback)
      {
        mrpt::tfest::TPotentialMatch pm;
        pm.idx_this = in_correspondences[idx].globalIdx;
        pm.idx_other = in_correspondences[idx].localIdx;
        if (!params.user_individual_compat_callback(pm)) continue;  // Skip this one!
      }

      // Consensus set: Maybe inliers + new point
      CPose3DQuat csOutQuat;
      mbInliers.push_back(in_correspondences[idx]);  // Insert
      res = mrpt::tfest::se3_l2(mbInliers, csOutQuat, scale, params.forceScaleToUnity);
      mbInliers.erase(mbInliers.end() - 1);  // Erase

      if (!res)
      {
        std::cerr << "[tfest::se3_l2_robust] tfest::se3_l2() returned "
                     "false for tentative subset during RANSAC "
                     "iteration!\n";
        continue;
      }

      // Is this point a supporter of the initial inlier group?
      const CPose3D csOut = CPose3D(csOutQuat);

      const double linDist = mbOut.distanceTo(csOut);
      const double angDist =
          mrpt::poses::Lie::SO<3>::log((csOut - mbOut).getRotationMatrix()).norm();
      const double scaleDist = std::abs(mbOut_vec[6] - scale);

      if (linDist < params.ransac_threshold_lin && angDist < params.ransac_threshold_ang &&
          scaleDist < params.ransac_threshold_scale)
      {
        // Inlier detected -> add to the inlier list
        cSet.push_back(idx);
      }
      else
      {
        // cout << " It " << iterations << " - RANSAC Outlier Detected:
        // " << k << endl;
      }
    }  // end 'inner' for

    // Test cSet size
    if (cSet.size() >= d)
    {
      // Good set of points found
      TMatchingPairList cSetInliers;
      cSetInliers.resize(cSet.size());
      for (size_t m = 0; m < cSet.size(); m++) cSetInliers[m] = in_correspondences[cSet[m]];

      // Compute output: Consensus Set + Initial Inliers Guess
      CPose3DQuat cIOutQuat;
      res = mrpt::tfest::se3_l2(
          cSetInliers, cIOutQuat, scale,
          params.forceScaleToUnity);  // Compute output
      ASSERTMSG_(
          res,
          "tfest::se3_l2() returned false for tentative subset during "
          "RANSAC iteration!");

      // Compute error for consensus_set
      const CPose3D cIOut = CPose3D(cIOutQuat);
      const double err = std::sqrt(
          square(mbOut_vec[0] - cIOut.x()) + square(mbOut_vec[1] - cIOut.y()) +
          square(mbOut_vec[2] - cIOut.z()) + square(mbOut_vec[3] - cIOut.yaw()) +
          square(mbOut_vec[4] - cIOut.pitch()) + square(mbOut_vec[5] - cIOut.roll()) +
          square(mbOut_vec[6] - scale));

      // Is the best set of points so far?
      if (err < min_err && cSet.size() >= max_size)
      {
        min_err = err;
        max_size = cSet.size();
        results.transformation = cIOutQuat;
        results.scale = scale;
        results.inliers_idx = cSet;
      }
      // printf(" - Consensus set size: %u - Error: %.6f\n", (unsigned
      // int)cSet.size(), err );
    }  // end if cSet.size() > d
    else
    {
      // printf(" - Consensus set size: %u - Not big enough!\n", (unsigned
      // int)cSet.size() );
    }
  }  // end 'iterations' for

  if (max_size == 0)
  {
    if (params.verbose) std::cerr << "[se3_l2_robust] No solution found, maximum size is == 0!\n";
    return false;
  }

  return true;
  MRPT_END
}
