// Copyright (c), ETH Zurich and UNC Chapel Hill.
// All rights reserved.
//
// Redistribution and use in source and binary forms, with or without
// modification, are permitted provided that the following conditions are met:
//
//     * Redistributions of source code must retain the above copyright
//       notice, this list of conditions and the following disclaimer.
//
//     * Redistributions in binary form must reproduce the above copyright
//       notice, this list of conditions and the following disclaimer in the
//       documentation and/or other materials provided with the distribution.
//
//     * Neither the name of ETH Zurich and UNC Chapel Hill nor the names of
//       its contributors may be used to endorse or promote products derived
//       from this software without specific prior written permission.
//
// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
// AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
// IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
// ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDERS OR CONTRIBUTORS BE
// LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
// CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
// SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
// INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
// CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
// ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
// POSSIBILITY OF SUCH DAMAGE.

#pragma once

#include "colmap/optim/random_sampler.h"
#include "colmap/optim/ransac.h"
#include "colmap/optim/support_measurement.h"
#include "colmap/util/logging.h"
#include "colmap/util/threading.h"

#include <atomic>
#include <cfloat>
#include <optional>
#include <vector>

#ifdef _OPENMP
#include <omp.h>
#endif

namespace colmap {

// Implementation of LO-RANSAC (Locally Optimized RANSAC).
//
// "Locally Optimized RANSAC" Ondrej Chum, Jiri Matas, Josef Kittler, DAGM 2003.
template <typename Estimator,
          typename LocalEstimator,
          typename SupportMeasurer = InlierSupportMeasurer,
          typename Sampler = RandomSampler>
class LORANSAC : public RANSAC<Estimator, SupportMeasurer, Sampler> {
 public:
  using typename RANSAC<Estimator, SupportMeasurer, Sampler>::Report;

  explicit LORANSAC(const RANSACOptions& options,
                    Estimator estimator = Estimator(),
                    LocalEstimator local_estimator = LocalEstimator(),
                    SupportMeasurer support_measurer = SupportMeasurer(),
                    Sampler sampler = Sampler(Estimator::kMinNumSamples));

  // Robustly estimate model with RANSAC (RANdom SAmple Consensus).
  //
  // @param X              Independent variables.
  // @param Y              Dependent variables.
  //
  // @return               The report with the results of the estimation.
  Report Estimate(const std::vector<typename Estimator::X_t>& X,
                  const std::vector<typename Estimator::Y_t>& Y);

  using RANSAC<Estimator, SupportMeasurer, Sampler>::estimator;
  using RANSAC<Estimator, SupportMeasurer, Sampler>::support_measurer;
  using RANSAC<Estimator, SupportMeasurer, Sampler>::sampler;
  LocalEstimator local_estimator;

 private:
  using RANSAC<Estimator, SupportMeasurer, Sampler>::options_;
};

////////////////////////////////////////////////////////////////////////////////
// Implementation
////////////////////////////////////////////////////////////////////////////////

template <typename Estimator,
          typename LocalEstimator,
          typename SupportMeasurer,
          typename Sampler>
LORANSAC<Estimator, LocalEstimator, SupportMeasurer, Sampler>::LORANSAC(
    const RANSACOptions& options,
    Estimator estimator,
    LocalEstimator local_estimator,
    SupportMeasurer support_measurer,
    Sampler sampler)
    : RANSAC<Estimator, SupportMeasurer, Sampler>(options,
                                                  std::move(estimator),
                                                  std::move(support_measurer),
                                                  std::move(sampler)),
      local_estimator(std::move(local_estimator)) {}

template <typename Estimator,
          typename LocalEstimator,
          typename SupportMeasurer,
          typename Sampler>
typename LORANSAC<Estimator, LocalEstimator, SupportMeasurer, Sampler>::Report
LORANSAC<Estimator, LocalEstimator, SupportMeasurer, Sampler>::Estimate(
    const std::vector<typename Estimator::X_t>& X,
    const std::vector<typename Estimator::Y_t>& Y) {
  THROW_CHECK_EQ(X.size(), Y.size());

  const size_t num_samples = X.size();

  typename RANSAC<Estimator, SupportMeasurer, Sampler>::Report report;
  report.success = false;
  report.num_trials = 0;

  if (num_samples < Estimator::kMinNumSamples) {
    return report;
  }

  const double max_residual = options_.max_error * options_.max_error;
  const size_t min_num_trials = options_.min_num_trials;

  sampler.Initialize(num_samples);
  const size_t max_num_trials =
      std::min<size_t>(options_.max_num_trials, sampler.MaxNumSamples());

  typename SupportMeasurer::Support best_support;
  std::optional<typename Estimator::M_t> best_model;
  bool best_model_is_local = false;

  [[maybe_unused]] const int num_threads =
      GetEffectiveNumThreads(options_.num_threads);
  if constexpr (!std::is_same_v<Sampler, RandomSampler>) {
    THROW_CHECK_EQ(num_threads, 1)
        << "Parallel LORANSAC only supports RandomSampler";
  }

  std::atomic<size_t> trial_counter(0);
  std::atomic<size_t> dyn_max_num_trials(max_num_trials);
  std::atomic<bool> abort_flag(false);

// This creates a parallel region that runs with num_threads threads (or
// serially when num_threads == 1 via the if-clause). Without OpenMP, the
// pragma is absent and the block runs once serially.
#ifdef _OPENMP
#pragma omp parallel num_threads(num_threads) if (num_threads > 1)
#endif
  {
    // Per-thread copies of mutable objects. Each thread needs its own sampler
    // (see random_sampler.h) and its own estimator/support_measurer instances.
    Sampler thread_sampler(Estimator::kMinNumSamples);
    thread_sampler.Initialize(num_samples);
    Estimator thread_estimator = estimator;
    LocalEstimator thread_local_estimator = local_estimator;
    SupportMeasurer thread_support_measurer = support_measurer;

    // Seed per-thread PRNG with distinct seed.
    if constexpr (is_randomized_sampler<Sampler>::value) {
      if (options_.random_seed != -1) {
#ifdef _OPENMP
        SetPRNGSeed(options_.random_seed + omp_get_thread_num());
#else
        SetPRNGSeed(options_.random_seed);
#endif
      }
    }

    // Per-thread working buffers.
    std::vector<double> residuals;
    std::vector<double> best_local_residuals;

    std::vector<typename LocalEstimator::X_t> X_inlier;
    std::vector<typename LocalEstimator::Y_t> Y_inlier;

    std::vector<typename Estimator::X_t> X_rand(Estimator::kMinNumSamples);
    std::vector<typename Estimator::Y_t> Y_rand(Estimator::kMinNumSamples);
    std::vector<typename Estimator::M_t> sample_models;
    std::vector<typename LocalEstimator::M_t> local_models;

    while (true) {
      const size_t curr_thread_trial =
          trial_counter.fetch_add(1, std::memory_order_relaxed);
      if (curr_thread_trial >= max_num_trials ||
          abort_flag.load(std::memory_order_relaxed)) {
        break;
      }

      thread_sampler.SampleXY(X, Y, &X_rand, &Y_rand);

      // Estimate model for current subset.
      sample_models.clear();
      thread_estimator.Estimate(X_rand, Y_rand, &sample_models);

      // Iterate through all estimated models.
      for (const auto& sample_model : sample_models) {
        thread_estimator.Residuals(X, Y, sample_model, &residuals);
        THROW_CHECK_EQ(residuals.size(), num_samples);

        const auto support =
            thread_support_measurer.Evaluate(residuals, max_residual);

        // Speculatively check if better than global best. In the parallel
        // case, multiple threads may pass this check concurrently — the
        // authoritative update below re-checks under the same lock.
        bool is_better = false;
#ifdef _OPENMP
#pragma omp critical
#endif
        {
          is_better =
              thread_support_measurer.IsLeftBetter(support, best_support);
        }

        if (is_better) {
          // Do local optimization (per-thread, outside critical section).
          typename SupportMeasurer::Support local_best_support = support;
          std::optional<typename Estimator::M_t> local_best_model =
              sample_model;
          bool local_best_is_local = false;

          // Estimate locally optimized model from inliers.
          if (support.num_inliers > Estimator::kMinNumSamples &&
              support.num_inliers >= LocalEstimator::kMinNumSamples) {
            // Recursive local optimization to expand inlier set.
            const size_t kMaxNumLocalTrials = 10;
            for (size_t local_num_trials = 0;
                 local_num_trials < kMaxNumLocalTrials;
                 ++local_num_trials) {
              X_inlier.clear();
              Y_inlier.clear();
              X_inlier.reserve(num_samples);
              Y_inlier.reserve(num_samples);
              for (size_t i = 0; i < residuals.size(); ++i) {
                if (residuals[i] <= max_residual) {
                  X_inlier.push_back(X[i]);
                  Y_inlier.push_back(Y[i]);
                }
              }

              local_models.clear();
              thread_local_estimator.Estimate(
                  X_inlier, Y_inlier, &local_models);

              const size_t prev_best_num_inliers =
                  local_best_support.num_inliers;

              for (const auto& local_model : local_models) {
                thread_local_estimator.Residuals(X, Y, local_model, &residuals);
                THROW_CHECK_EQ(residuals.size(), num_samples);

                const auto local_support =
                    thread_support_measurer.Evaluate(residuals, max_residual);

                // Check if locally optimized model is better.
                if (thread_support_measurer.IsLeftBetter(local_support,
                                                         local_best_support)) {
                  local_best_support = local_support;
                  local_best_model = local_model;
                  local_best_is_local = true;
                  std::swap(residuals, best_local_residuals);
                }
              }

              // Only continue recursive local optimization, if the inlier set
              // size increased and we thus have a chance to further improve.
              if (local_best_support.num_inliers <= prev_best_num_inliers) {
                break;
              }

              // Swap back the residuals, so we can extract the best inlier
              // set in the next recursion of local optimization.
              std::swap(residuals, best_local_residuals);
            }
          }

          // Commit local optimization result to global best under lock.
#ifdef _OPENMP
#pragma omp critical
#endif
          {
            if (thread_support_measurer.IsLeftBetter(local_best_support,
                                                     best_support)) {
              best_support = local_best_support;
              best_model = local_best_model;
              best_model_is_local = local_best_is_local;

              dyn_max_num_trials.store(
                  RANSAC<Estimator, SupportMeasurer, Sampler>::ComputeNumTrials(
                      best_support.num_inliers,
                      num_samples,
                      options_.confidence,
                      options_.dyn_num_trials_multiplier),
                  std::memory_order_relaxed);
            }
          }
        }

        if (curr_thread_trial >=
                dyn_max_num_trials.load(std::memory_order_relaxed) &&
            curr_thread_trial >= min_num_trials) {
          abort_flag.store(true, std::memory_order_relaxed);
          break;
        }
      }

      if (abort_flag.load(std::memory_order_relaxed)) {
        break;
      }
    }
  }

  report.num_trials = trial_counter.load();

  if (!best_model.has_value()) {
    return report;
  }

  report.support = best_support;
  report.model = best_model.value();

  // No valid model was found.
  if (report.support.num_inliers < estimator.kMinNumSamples) {
    return report;
  }

  report.success = true;

  // Determine inlier mask. Note that this calculates the residuals for the
  // best model twice, but saves to copy and fill the inlier mask for each
  // evaluated model. Some benchmarking revealed that this approach is faster.

  std::vector<double> residuals(num_samples);
  if (best_model_is_local) {
    local_estimator.Residuals(X, Y, report.model, &residuals);
  } else {
    estimator.Residuals(X, Y, report.model, &residuals);
  }

  THROW_CHECK_EQ(residuals.size(), num_samples);

  report.inlier_mask.resize(num_samples);
  for (size_t i = 0; i < residuals.size(); ++i) {
    report.inlier_mask[i] = residuals[i] <= max_residual;
  }

  return report;
}

}  // namespace colmap
