// 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.

#include "colmap/controllers/incremental_pipeline.h"

#include "colmap/estimators/alignment.h"
#include "colmap/estimators/bundle_adjustment_caspar.h"
#include "colmap/estimators/bundle_adjustment_ceres.h"
#include "colmap/scene/database.h"
#include "colmap/util/file.h"
#include "colmap/util/timer.h"

namespace colmap {
namespace {

void CustomizeIncrementalPipelineOptions(const DatabaseCache& database_cache,
                                         IncrementalPipelineOptions& options) {
  // If the total number of images is small then do not enforce the
  // minimum model size so that we can reconstruct small image
  // collections, i.e., if the model is at least half of the total number
  // of images, we always keep it.
  options.min_model_size = std::min<size_t>(0.5 * database_cache.NumImages(),
                                            options.min_model_size);
}

DatabaseCache::Options CreateDatabaseCacheOptions(
    const IncrementalPipelineOptions& options,
    const ReconstructionManager& reconstruction_manager) {
  DatabaseCache::Options database_cache_options;
  database_cache_options.min_num_matches =
      static_cast<size_t>(options.min_num_matches);
  database_cache_options.ignore_watermarks = options.ignore_watermarks;
  database_cache_options.image_names = {options.image_names.begin(),
                                        options.image_names.end()};
  // Make sure images of the given reconstruction are also included when
  // manually specifying images for the reconstruction procedure.
  if (reconstruction_manager.Size() == 1 && !options.image_names.empty()) {
    const auto& reconstruction = reconstruction_manager.Get(0);
    for (const image_t image_id : reconstruction->RegImageIds()) {
      const auto& image = reconstruction->Image(image_id);
      database_cache_options.image_names.insert(image.Name());
    }
  }
  database_cache_options.load_all_images = options.load_all_images;
  database_cache_options.convert_pose_priors_to_enu =
      options.use_prior_position;
  return database_cache_options;
}

void IterativeGlobalRefinement(const IncrementalPipelineOptions& options,
                               const IncrementalMapper::Options& mapper_options,
                               IncrementalMapper& mapper) {
  LOG(INFO) << "Retriangulation and Global bundle adjustment";
  mapper.IterativeGlobalRefinement(options.ba_global_max_refinements,
                                   options.ba_global_max_refinement_change,
                                   mapper_options,
                                   options.GlobalBundleAdjustment(),
                                   options.Triangulation());
  mapper.FilterFrames(mapper_options);
}

void ExtractColors(const std::filesystem::path& image_path,
                   const image_t image_id,
                   Reconstruction& reconstruction) {
  if (!reconstruction.ExtractColorsForImage(image_id, image_path)) {
    LOG(WARNING) << "Could not read image "
                 << reconstruction.Image(image_id).Name() << " at path "
                 << image_path << ".";
  }
}

void WriteSnapshot(const Reconstruction& reconstruction,
                   const std::filesystem::path& snapshot_path) {
  LOG(INFO) << "Creating snapshot";
  // Get the current timestamp in milliseconds.
  const size_t timestamp =
      std::chrono::duration_cast<std::chrono::milliseconds>(
          std::chrono::high_resolution_clock::now().time_since_epoch())
          .count();
  // Write reconstruction to unique path with current timestamp.
  const auto path = snapshot_path / StringPrintf("%010zu", timestamp);
  CreateDirIfNotExists(path);
  VLOG(1) << "=> Writing to " << path;
  reconstruction.Write(path);
}

bool HasUnknownSensorFromRig(const Reconstruction& reconstruction) {
  std::unordered_set<const Rig*> parameterized_rigs;
  for (const auto& [_, image] : reconstruction.Images()) {
    parameterized_rigs.insert(image.FramePtr()->RigPtr());
  }
  for (const Rig* rig : parameterized_rigs) {
    for (const auto& [sensor_id, sensor_from_rig] : rig->NonRefSensors()) {
      if (sensor_id.type == SensorType::CAMERA &&
          !sensor_from_rig.has_value()) {
        return true;
      }
    }
  }
  return false;
}

}  // namespace

IncrementalMapper::Options IncrementalPipelineOptions::Mapper() const {
  IncrementalMapper::Options options = mapper;
  options.abs_pose_refine_focal_length = ba_refine_focal_length;
  options.abs_pose_refine_extra_params = ba_refine_extra_params;
  options.min_focal_length_ratio = min_focal_length_ratio;
  options.max_focal_length_ratio = max_focal_length_ratio;
  options.max_extra_param = max_extra_param;
  options.num_threads = num_threads;
  options.fix_existing_frames = fix_existing_frames;
  options.constant_rigs = constant_rigs;
  options.constant_cameras = constant_cameras;
  options.use_prior_position = use_prior_position;
  options.use_robust_loss_on_prior_position = use_robust_loss_on_prior_position;
  options.prior_position_loss_scale = prior_position_loss_scale;
  options.random_seed = random_seed;
  return options;
}

IncrementalTriangulator::Options IncrementalPipelineOptions::Triangulation()
    const {
  IncrementalTriangulator::Options options = triangulation;
  options.min_focal_length_ratio = min_focal_length_ratio;
  options.max_focal_length_ratio = max_focal_length_ratio;
  options.max_extra_param = max_extra_param;
  options.random_seed = random_seed;
  return options;
}

BundleAdjustmentOptions IncrementalPipelineOptions::LocalBundleAdjustment()
    const {
  BundleAdjustmentOptions options;
  options.print_summary = false;
  options.backend = ba_local_backend;
  options.refine_focal_length = ba_refine_focal_length;
  options.refine_principal_point = ba_refine_principal_point;
  options.refine_extra_params = ba_refine_extra_params;
  options.refine_sensor_from_rig = ba_refine_sensor_from_rig;
  if (options.ceres) {
    options.ceres->solver_options.function_tolerance =
        ba_local_function_tolerance;
    options.ceres->solver_options.gradient_tolerance = 10.0;
    options.ceres->solver_options.parameter_tolerance = 0.0;
    options.ceres->solver_options.max_num_iterations =
        ba_local_max_num_iterations;
    options.ceres->solver_options.max_linear_solver_iterations = 100;
    options.ceres->solver_options.logging_type = ceres::LoggingType::SILENT;
    options.ceres->solver_options.num_threads = num_threads;
#if CERES_VERSION_MAJOR < 2
    options.ceres->solver_options.num_linear_solver_threads = num_threads;
#endif  // CERES_VERSION_MAJOR
    options.ceres->min_num_residuals_for_cpu_multi_threading =
        ba_min_num_residuals_for_cpu_multi_threading;
    options.ceres->loss_function_scale = 1.0;
    options.ceres->loss_function_type =
        CeresBundleAdjustmentOptions::LossFunctionType::SOFT_L1;
    options.ceres->use_gpu = ba_use_gpu;
    options.ceres->gpu_index = ba_gpu_index;
  }
  if (options.caspar) {
    options.caspar->gpu_index = ba_gpu_index;
  }
  return options;
}

BundleAdjustmentOptions IncrementalPipelineOptions::GlobalBundleAdjustment()
    const {
  BundleAdjustmentOptions options;
  options.print_summary = false;
  options.backend = ba_global_backend;
  options.refine_focal_length = ba_refine_focal_length;
  options.refine_principal_point = ba_refine_principal_point;
  options.refine_extra_params = ba_refine_extra_params;
  options.refine_sensor_from_rig = ba_refine_sensor_from_rig;
  if (options.ceres) {
    options.ceres->solver_options.function_tolerance =
        ba_global_function_tolerance;
    options.ceres->solver_options.gradient_tolerance = 1.0;
    options.ceres->solver_options.parameter_tolerance = 0.0;
    options.ceres->solver_options.max_num_iterations =
        ba_global_max_num_iterations;
    options.ceres->solver_options.max_linear_solver_iterations = 100;
    options.ceres->solver_options.logging_type = ceres::LoggingType::SILENT;
    if (VLOG_IS_ON(2)) {
      options.ceres->solver_options.minimizer_progress_to_stdout = true;
      options.ceres->solver_options.logging_type =
          ceres::LoggingType::PER_MINIMIZER_ITERATION;
    }
    options.ceres->solver_options.num_threads = num_threads;
#if CERES_VERSION_MAJOR < 2
    options.ceres->solver_options.num_linear_solver_threads = num_threads;
#endif  // CERES_VERSION_MAJOR
    options.ceres->min_num_residuals_for_cpu_multi_threading =
        ba_min_num_residuals_for_cpu_multi_threading;
    options.ceres->loss_function_type =
        CeresBundleAdjustmentOptions::LossFunctionType::TRIVIAL;
    options.ceres->use_gpu = ba_use_gpu;
    options.ceres->gpu_index = ba_gpu_index;
  }
  if (options.caspar) {
    options.caspar->gpu_index = ba_gpu_index;
  }
  return options;
}

bool IncrementalPipelineOptions::Check() const {
  CHECK_OPTION_GT(min_num_matches, 0);
  CHECK_OPTION_GT(max_num_models, 0);
  CHECK_OPTION_GT(max_model_overlap, 0);
  CHECK_OPTION_GE(min_model_size, 0);
  CHECK_OPTION_GT(init_num_trials, 0);
  CHECK_OPTION_GT(min_focal_length_ratio, 0);
  CHECK_OPTION_GT(max_focal_length_ratio, 0);
  CHECK_OPTION_GE(max_extra_param, 0);
  CHECK_OPTION_GE(ba_local_max_num_iterations, 0);
  CHECK_OPTION_GT(ba_global_frames_ratio, 1.0);
  CHECK_OPTION_GT(ba_global_points_ratio, 1.0);
  CHECK_OPTION_GT(ba_global_frames_freq, 0);
  CHECK_OPTION_GT(ba_global_points_freq, 0);
  CHECK_OPTION_GT(ba_global_max_num_iterations, 0);
  CHECK_OPTION_GT(ba_local_max_refinements, 0);
  CHECK_OPTION_GE(ba_local_max_refinement_change, 0);
  CHECK_OPTION_GE(ba_global_max_refinements, 0);
  CHECK_OPTION_GE(ba_global_max_refinement_change, 0);
  CHECK_OPTION_GE(snapshot_frames_freq, 0);
  CHECK_OPTION_GT(prior_position_loss_scale, 0.);
  CHECK_OPTION_GE(num_threads, -1);
  CHECK_OPTION_GE(random_seed, -1);
#ifndef CASPAR_ENABLED
  CHECK_OPTION(ba_local_backend != BundleAdjustmentBackend::CASPAR);
  CHECK_OPTION(ba_global_backend != BundleAdjustmentBackend::CASPAR);
#endif
  CHECK_OPTION(Mapper().Check());
  CHECK_OPTION(Triangulation().Check());
  return true;
}

IncrementalPipeline::IncrementalPipeline(
    std::shared_ptr<IncrementalPipelineOptions> options,
    std::shared_ptr<class Database> database,
    std::shared_ptr<class ReconstructionManager> reconstruction_manager)
    : options_(std::move(THROW_CHECK_NOTNULL(options))),
      reconstruction_manager_(
          THROW_CHECK_NOTNULL(std::move(reconstruction_manager))),
      total_run_timer_(std::make_shared<Timer>()) {
  THROW_CHECK(options_->Check());
  THROW_CHECK_NOTNULL(database);

  LOG(INFO) << "Loading database";
  Timer timer;
  timer.Start();
  database_cache_ = DatabaseCache::Create(
      *database,
      CreateDatabaseCacheOptions(*options_, *reconstruction_manager_));
  timer.PrintMinutes();

  CustomizeIncrementalPipelineOptions(*database_cache_, *options_);

  RegisterCallbacks();
}

IncrementalPipeline::IncrementalPipeline(
    std::shared_ptr<IncrementalPipelineOptions> options,
    std::shared_ptr<class DatabaseCache> database_cache,
    std::shared_ptr<class ReconstructionManager> reconstruction_manager)
    : options_(std::move(THROW_CHECK_NOTNULL(options))),
      reconstruction_manager_(
          THROW_CHECK_NOTNULL(std::move(reconstruction_manager))),
      total_run_timer_(std::make_shared<Timer>()) {
  THROW_CHECK(options_->Check());
  THROW_CHECK_NOTNULL(database_cache);

  database_cache_ = DatabaseCache::CreateFromCache(
      *database_cache,
      CreateDatabaseCacheOptions(*options_, *reconstruction_manager_));

  CustomizeIncrementalPipelineOptions(*database_cache_, *options_);

  RegisterCallbacks();
}

void IncrementalPipeline::Run() {
  total_run_timer_->Start();

  if (database_cache_->NumImages() == 0) {
    LOG(WARNING) << "No images with matches";
    return;
  }

  if (options_->use_prior_position && database_cache_->NumPosePriors() == 0) {
    LOG(WARNING) << "No pose priors";
    return;
  }

  // Is there a sub-reconstruction before we start the reconstruction? I.e. the
  // user has imported an existing reconstruction.
  const bool continue_reconstruction = reconstruction_manager_->Size() > 0;
  THROW_CHECK_LE(reconstruction_manager_->Size(), 1)
      << "Can only continue from a single reconstruction, "
         "but multiple are given.";

  const size_t num_images = database_cache_->NumImages();

  IncrementalMapper::Options mapper_options = options_->Mapper();
  IncrementalMapper mapper(database_cache_);
  if (Reconstruct(mapper,
                  mapper_options,
                  /*continue_reconstruction=*/continue_reconstruction) ==
      Status::STOP) {
    total_run_timer_->PrintMinutes();
    return;
  }

  auto ShouldStop = [this, &mapper, &num_images]() {
    return mapper.NumTotalRegImages() == num_images || CheckIfStopped() ||
           CheckReachedMaxRuntime();
  };

  const size_t kNumInitRelaxations = 2;
  for (size_t i = 0; i < kNumInitRelaxations; ++i) {
    if (ShouldStop()) {
      break;
    }

    LOG(INFO) << "=> Relaxing the initialization constraints.";
    mapper_options.init_min_num_inliers /= 2;
    mapper.ResetInitializationStats();
    if (Reconstruct(mapper,
                    mapper_options,
                    /*continue_reconstruction=*/false) == Status::STOP) {
      break;
    }

    if (ShouldStop()) {
      break;
    }

    LOG(INFO) << "=> Relaxing the initialization constraints.";
    mapper_options.init_min_tri_angle /= 2;
    mapper.ResetInitializationStats();
    if (Reconstruct(mapper,
                    mapper_options,
                    /*continue_reconstruction=*/false) == Status::STOP) {
      break;
    }
  }

  total_run_timer_->PrintMinutes();
}

IncrementalPipeline::Status IncrementalPipeline::InitializeReconstruction(
    IncrementalMapper& mapper,
    const IncrementalMapper::Options& mapper_options,
    Reconstruction& reconstruction) {
  image_t image_id1 = static_cast<image_t>(options_->init_image_id1);
  image_t image_id2 = static_cast<image_t>(options_->init_image_id2);

  // Try to find good initial pair.
  Rigid3d cam2_from_cam1;
  if (!options_->IsInitialPairProvided()) {
    LOG(INFO) << "Finding good initial image pair";
    const bool find_init_success = mapper.FindInitialImagePair(
        mapper_options, image_id1, image_id2, cam2_from_cam1);
    if (!find_init_success) {
      LOG(INFO) << "=> No good initial image pair found.";
      return Status::NO_INITIAL_PAIR;
    }
  } else {
    if (!reconstruction.ExistsImage(image_id1) ||
        !reconstruction.ExistsImage(image_id2)) {
      LOG(INFO) << StringPrintf(
          "=> Initial image pair #%d and #%d does not exist.",
          image_id1,
          image_id2);
      return Status::NO_INITIAL_PAIR;
    }
    const bool provided_init_success = mapper.EstimateInitialTwoViewGeometry(
        mapper_options, image_id1, image_id2, cam2_from_cam1);
    if (!provided_init_success) {
      LOG(INFO) << "=> Provided pair is unsuitable for initialization.";
      return Status::BAD_INITIAL_PAIR;
    }
  }

  LOG(INFO) << StringPrintf(
      "Registering initial image pair #%d and #%d", image_id1, image_id2);
  mapper.RegisterInitialImagePair(
      mapper_options, image_id1, image_id2, cam2_from_cam1);

  IncrementalTriangulator::Options tri_options = options_->Triangulation();
  tri_options.min_angle = mapper_options.init_min_tri_angle;
  for (const image_t image_id : {image_id1, image_id2}) {
    const Image& image = reconstruction.Image(image_id);
    for (const data_t& data_id : image.FramePtr()->ImageIds()) {
      mapper.TriangulateImage(tri_options, data_id.id);
    }
  }

  if (reconstruction.NumPoints3D() == 0) {
    return Status::BAD_INITIAL_PAIR;
  }

  LOG(INFO) << "Global bundle adjustment";
  mapper.AdjustGlobalBundle(mapper_options, options_->GlobalBundleAdjustment());
  reconstruction.Normalize();
  mapper.FilterPoints(mapper_options);
  mapper.FilterFrames(mapper_options);

  // Initial image pair failed to register.
  if (reconstruction.NumRegFrames() == 0 || reconstruction.NumPoints3D() == 0) {
    return Status::BAD_INITIAL_PAIR;
  }

  // Number of triangulated points not enough for registering future images.
  if (static_cast<int>(reconstruction.NumPoints3D()) <
      mapper_options.abs_pose_min_num_inliers) {
    return Status::BAD_INITIAL_PAIR;
  }

  if (options_->extract_colors) {
    for (const image_t image_id : {image_id1, image_id2}) {
      const Image& image = reconstruction.Image(image_id);
      for (const data_t& data_id : image.FramePtr()->ImageIds()) {
        ExtractColors(options_->image_path, data_id.id, reconstruction);
      }
    }
  }
  return Status::SUCCESS;
}

bool IncrementalPipeline::CheckRunGlobalRefinement(
    const Reconstruction& reconstruction,
    const size_t ba_prev_num_reg_frames,
    const size_t ba_prev_num_points) {
  return reconstruction.NumRegFrames() >=
             options_->ba_global_frames_ratio * ba_prev_num_reg_frames ||
         reconstruction.NumRegFrames() >=
             options_->ba_global_frames_freq + ba_prev_num_reg_frames ||
         reconstruction.NumPoints3D() >=
             options_->ba_global_points_ratio * ba_prev_num_points ||
         reconstruction.NumPoints3D() >=
             options_->ba_global_points_freq + ba_prev_num_points;
}

IncrementalPipeline::Status IncrementalPipeline::ReconstructSubModel(
    IncrementalMapper& mapper,
    const IncrementalMapper::Options& mapper_options,
    const std::shared_ptr<Reconstruction>& reconstruction) {
  mapper.BeginReconstruction(reconstruction);

  if (HasUnknownSensorFromRig(*reconstruction)) {
    return Status::UNKNOWN_SENSOR_FROM_RIG;
  }

  ////////////////////////////////////////////////////////////////////////////
  // Register initial pair
  ////////////////////////////////////////////////////////////////////////////

  if (reconstruction->NumRegFrames() == 0) {
    const Status init_status = IncrementalPipeline::InitializeReconstruction(
        mapper, mapper_options, *reconstruction);
    if (init_status != Status::SUCCESS) {
      return init_status;
    }
  }
  Callback(INITIAL_IMAGE_PAIR_REG_CALLBACK);

  ////////////////////////////////////////////////////////////////////////////
  // Incremental mapping
  ////////////////////////////////////////////////////////////////////////////

  size_t snapshot_prev_num_reg_frames = reconstruction->NumRegFrames();
  size_t ba_prev_num_reg_frames = reconstruction->NumRegFrames();
  size_t ba_prev_num_points = reconstruction->NumPoints3D();

  std::vector<bool> structure_less_flags;
  if (options_->structure_less_registration_only) {
    structure_less_flags = {true};
  } else {
    if (options_->structure_less_registration_fallback) {
      structure_less_flags = {false, true};
    } else {
      structure_less_flags = {false};
    }
  }

  bool reg_next_success = true;
  bool prev_reg_next_success = true;
  do {
    if (CheckIfStopped() || CheckReachedMaxRuntime()) {
      break;
    }

    prev_reg_next_success = reg_next_success;
    reg_next_success = false;
    image_t next_image_id = kInvalidImageId;

    // Try to register next image. Always prefer structure-based registration
    // first, and if that fails, try (less reliable) structure-less
    // registration.
    for (const bool structure_less : structure_less_flags) {
      const std::vector<image_t> next_images = mapper.FindNextImages(
          mapper_options, /*structure_less=*/structure_less);

      for (size_t reg_trial = 0; reg_trial < next_images.size(); ++reg_trial) {
        next_image_id = next_images[reg_trial];

        LOG(INFO) << StringPrintf("Registering image #%d (num_reg_frames=%d)",
                                  next_image_id,
                                  reconstruction->NumRegFrames());
        LOG(INFO) << StringPrintf(
            "=> Image sees %d / %d points",
            mapper.ObservationManager().NumVisiblePoints3D(next_image_id),
            mapper.ObservationManager().NumObservations(next_image_id));

        if (structure_less) {
          LOG(INFO) << StringPrintf(
              "Registering image with structure-less fallback");
          LOG(INFO) << StringPrintf(
              "=> Image sees %d / %d correspondences",
              mapper.ObservationManager().NumVisibleCorrespondences(
                  next_image_id),
              mapper.ObservationManager().NumCorrespondences(next_image_id));
          reg_next_success = mapper.RegisterNextStructureLessImage(
              mapper_options, next_image_id);
        } else {
          reg_next_success =
              mapper.RegisterNextImage(mapper_options, next_image_id);
        }

        if (reg_next_success) {
          break;
        } else {
          LOG(INFO) << "=> Could not register, trying another image.";

          // If initial model fails to continue for some time,
          // abort and try different initial pair.
          const size_t kMinNumInitialRegTrials = 30;
          if (reg_trial >= kMinNumInitialRegTrials &&
              reconstruction->NumRegImages() <
                  static_cast<size_t>(options_->min_model_size)) {
            break;
          }
        }
      }

      if (reg_next_success) {
        break;
      }
    }

    if (reg_next_success) {
      const Image& image = reconstruction->Image(next_image_id);
      for (const data_t& data_id : image.FramePtr()->ImageIds()) {
        mapper.TriangulateImage(options_->Triangulation(), data_id.id);
      }
      mapper.IterativeLocalRefinement(options_->ba_local_max_refinements,
                                      options_->ba_local_max_refinement_change,
                                      mapper_options,
                                      options_->LocalBundleAdjustment(),
                                      options_->Triangulation(),
                                      next_image_id);

      if (CheckRunGlobalRefinement(
              *reconstruction, ba_prev_num_reg_frames, ba_prev_num_points)) {
        IterativeGlobalRefinement(*options_, mapper_options, mapper);
        ba_prev_num_points = reconstruction->NumPoints3D();
        ba_prev_num_reg_frames = reconstruction->NumRegFrames();
      }

      if (options_->extract_colors) {
        for (const data_t& data_id : image.FramePtr()->ImageIds()) {
          ExtractColors(options_->image_path, data_id.id, *reconstruction);
        }
      }

      if (options_->snapshot_frames_freq > 0 &&
          reconstruction->NumRegFrames() >=
              options_->snapshot_frames_freq + snapshot_prev_num_reg_frames) {
        snapshot_prev_num_reg_frames = reconstruction->NumRegFrames();
        WriteSnapshot(*reconstruction, options_->snapshot_path);
      }

      Callback(NEXT_IMAGE_REG_CALLBACK);
    }

    const size_t max_model_overlap =
        static_cast<size_t>(options_->max_model_overlap);
    if (mapper.NumSharedRegImages() >= max_model_overlap) {
      break;
    }

    // If no image could be registered, try a single final global iterative
    // bundle adjustment and try again to register one image. If this fails
    // once, then exit the incremental mapping.
    if (!reg_next_success && prev_reg_next_success) {
      IterativeGlobalRefinement(*options_, mapper_options, mapper);
    }
  } while (reg_next_success || prev_reg_next_success);

  if (CheckIfStopped() || CheckReachedMaxRuntime()) {
    return Status::INTERRUPTED;
  }

  // Only run final global BA, if last incremental BA was not global.
  if (reconstruction->NumRegFrames() > 0 &&
      reconstruction->NumRegFrames() != ba_prev_num_reg_frames &&
      reconstruction->NumPoints3D() != ba_prev_num_points) {
    IterativeGlobalRefinement(*options_, mapper_options, mapper);
  }
  return Status::SUCCESS;
}

IncrementalPipeline::Status IncrementalPipeline::Reconstruct(
    IncrementalMapper& mapper,
    const IncrementalMapper::Options& mapper_options,
    bool continue_reconstruction) {
  for (int num_trials = 0; num_trials < options_->init_num_trials;
       ++num_trials) {
    if (CheckIfStopped() || CheckReachedMaxRuntime()) {
      return Status::STOP;
    }

    const size_t reconstruction_idx =
        (!continue_reconstruction || num_trials > 0)
            ? reconstruction_manager_->Add()
            : 0;
    std::shared_ptr<Reconstruction> reconstruction =
        reconstruction_manager_->Get(reconstruction_idx);

    const Status status =
        ReconstructSubModel(mapper, mapper_options, reconstruction);
    switch (status) {
      case Status::INTERRUPTED: {
        reconstruction->UpdatePoint3DErrors();
        LOG(INFO) << "Keeping reconstruction due to interrupt";
        mapper.EndReconstruction(/*discard=*/false);
        AlignReconstructionToOrigRigScales(database_cache_->Rigs(),
                                           reconstruction.get());
        return Status::STOP;
      }

      case Status::UNKNOWN_SENSOR_FROM_RIG: {
        LOG(ERROR)
            << "Discarding reconstruction due to unknown sensor_from_rig "
               "poses. Either explicitly define the poses by configuring the "
               "rigs or first run reconstruction without configured rigs and "
               "then derive the poses from the initial reconstruction for a "
               "subsequent reconstruction with rig constraints. See "
               "documentation for detailed instructions.";
        mapper.EndReconstruction(/*discard=*/true);
        reconstruction_manager_->Delete(reconstruction_idx);
        // If the reconstruction was discarded due to an unknown sensor from
        // rig, we can stop the outer trial loop, because all trials will fail.
        return Status::STOP;
      }

      case Status::BAD_INITIAL_PAIR: {
        LOG(INFO) << "Discarding reconstruction due to bad initial pair";
        mapper.EndReconstruction(/*discard=*/true);
        reconstruction_manager_->Delete(reconstruction_idx);
        // If an initial pair was found but it was bad, we discard and attempt
        // to initialize from any of the remaining pairs in the next trials.
        break;
      }

      case Status::NO_INITIAL_PAIR: {
        LOG(INFO) << "Discarding reconstruction due to no initial pair";
        mapper.EndReconstruction(/*discard=*/true);
        reconstruction_manager_->Delete(reconstruction_idx);
        // If no pair could be found, we can exit the trial loop, because
        // the next trials in this loop will not find anything unless the
        // initialization thresholds are relaxed. However, by relaxing the
        // constraints in the outer loop we can succeed.
        return Status::CONTINUE;
      }

      case Status::SUCCESS: {
        // Remember the total number of registered images before potentially
        // discarding it below due to small size, so we can exit out of the main
        // loop, if all images were registered.
        const size_t num_reg_images = reconstruction->NumRegImages();
        const size_t total_num_reg_images = mapper.NumTotalRegImages();

        // Always keep the first reconstruction, independent of size.
        if ((options_->multiple_models && reconstruction_manager_->Size() > 1 &&
             num_reg_images < static_cast<size_t>(options_->min_model_size)) ||
            num_reg_images == 0) {
          LOG(WARNING) << "Discarding reconstruction due to insufficient size";
          mapper.EndReconstruction(/*discard=*/true);
          reconstruction_manager_->Delete(reconstruction_idx);
        } else {
          reconstruction->UpdatePoint3DErrors();
          LOG(INFO) << "Keeping successful reconstruction";
          mapper.EndReconstruction(/*discard=*/false);
          AlignReconstructionToOrigRigScales(database_cache_->Rigs(),
                                             reconstruction.get());
        }

        Callback(LAST_IMAGE_REG_CALLBACK);

        // Check if we should or can reconstruct another sub-model.
        if (!options_->multiple_models ||
            reconstruction_manager_->Size() >=
                static_cast<size_t>(options_->max_num_models) ||
            total_num_reg_images >= database_cache_->NumImages() - 1) {
          return Status::STOP;
        }

        // In case the reconstruction was successful and there are remaining
        // images, we try to reconstruct another sub-model in the next trial.
        break;
      }

      default:
        LOG(FATAL_THROW) << "Unknown reconstruction status.";
    }
  }

  return Status::CONTINUE;
}

void IncrementalPipeline::TriangulateReconstruction(
    const std::shared_ptr<Reconstruction>& reconstruction) {
  THROW_CHECK_GT(database_cache_->NumImages(), 0)
      << "No images with matches found in the database";
  IncrementalMapper mapper(database_cache_);
  mapper.BeginReconstruction(reconstruction);

  LOG(INFO) << "Iterative triangulation";
  size_t image_idx = 0;
  for (const image_t image_id : reconstruction->RegImageIds()) {
    const auto& image = reconstruction->Image(image_id);

    LOG(INFO) << StringPrintf(
        "Triangulating image #%d (%d)", image_id, image_idx++);
    const size_t num_existing_points3D = image.NumPoints3D();
    LOG(INFO) << "=> Image sees " << num_existing_points3D << " / "
              << mapper.ObservationManager().NumObservations(image_id)
              << " points";

    mapper.TriangulateImage(options_->Triangulation(), image_id);
    VLOG(1) << "=> Triangulated "
            << (image.NumPoints3D() - num_existing_points3D) << " points";
  }

  LOG(INFO) << "Retriangulation and Global bundle adjustment";
  mapper.IterativeGlobalRefinement(options_->ba_global_max_refinements,
                                   options_->ba_global_max_refinement_change,
                                   options_->Mapper(),
                                   options_->GlobalBundleAdjustment(),
                                   options_->Triangulation(),
                                   /*normalize_reconstruction=*/false);
  mapper.EndReconstruction(/*discard=*/false);

  reconstruction->UpdatePoint3DErrors();

  LOG(INFO) << "Extracting colors";
  reconstruction->ExtractColorsForAllImages(options_->image_path,
                                            options_->num_threads);
}

void IncrementalPipeline::RegisterCallbacks() {
  RegisterCallback(INITIAL_IMAGE_PAIR_REG_CALLBACK);
  RegisterCallback(NEXT_IMAGE_REG_CALLBACK);
  RegisterCallback(LAST_IMAGE_REG_CALLBACK);
}

bool IncrementalPipeline::CheckReachedMaxRuntime() const {
  if (options_->max_runtime_seconds > 0 &&
      total_run_timer_->ElapsedSeconds() > options_->max_runtime_seconds) {
    LOG(INFO) << "Reached maximum runtime of " << options_->max_runtime_seconds
              << " seconds.";
    return true;
  }
  return false;
}

}  // namespace colmap
