// 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/scene/database_cache.h"

#include "colmap/geometry/gps.h"
#include "colmap/util/string.h"
#include "colmap/util/timer.h"

namespace colmap {
namespace {

bool UseInlierMatchesCheck(const DatabaseCache::Options& options,
                           int two_view_geometry_config,
                           size_t num_matches) {
  return num_matches >= options.min_num_matches &&
         (!options.ignore_watermarks ||
          two_view_geometry_config != TwoViewGeometry::WATERMARK);
};

std::vector<Eigen::Vector2d> FeatureKeypointsToPointsVector(
    const FeatureKeypoints& keypoints) {
  std::vector<Eigen::Vector2d> points(keypoints.size());
  for (size_t i = 0; i < keypoints.size(); ++i) {
    points[i] = Eigen::Vector2d(keypoints[i].x, keypoints[i].y);
  }
  return points;
}

}  // namespace

DatabaseCache::DatabaseCache()
    : correspondence_graph_(std::make_shared<class CorrespondenceGraph>()) {}

void DatabaseCache::Load(const Database& database, const Options& options) {
  const bool has_rigs = database.NumRigs() > 0;
  const bool has_frames = database.NumFrames() > 0;

  //////////////////////////////////////////////////////////////////////////////
  // Load rigs
  //////////////////////////////////////////////////////////////////////////////

  Timer timer;

  timer.Start();
  LOG(INFO) << "Loading rigs...";

  {
    std::vector<class Rig> rigs = database.ReadAllRigs();
    rigs_.reserve(rigs.size());
    for (auto& rig : rigs) {
      rigs_.emplace(rig.RigId(), std::move(rig));
    }
  }

  LOG(INFO) << StringPrintf(
      " %d in %.3fs", rigs_.size(), timer.ElapsedSeconds());

  //////////////////////////////////////////////////////////////////////////////
  // Load cameras
  //////////////////////////////////////////////////////////////////////////////

  timer.Restart();
  LOG(INFO) << "Loading cameras...";

  {
    std::vector<struct Camera> cameras = database.ReadAllCameras();
    cameras_.reserve(cameras.size());
    for (auto& camera : cameras) {
      if (!has_rigs) {
        // For backwards compatibility with old databases from before having
        // support for rigs/frames, we create a rig for each camera.
        class Rig rig;
        rig.SetRigId(camera.camera_id);
        rig.AddRefSensor(camera.SensorId());
        rigs_.emplace(rig.RigId(), std::move(rig));
      }
      cameras_.emplace(camera.camera_id, std::move(camera));
    }
  }

  LOG(INFO) << StringPrintf(
      " %d in %.3fs", cameras_.size(), timer.ElapsedSeconds());

  //////////////////////////////////////////////////////////////////////////////
  // Load frames
  //////////////////////////////////////////////////////////////////////////////

  timer.Restart();
  LOG(INFO) << "Loading frames...";

  {
    std::vector<class Frame> frames = database.ReadAllFrames();
    frames_.reserve(frames.size());
    for (auto& frame : frames) {
      frames_.emplace(frame.FrameId(), std::move(frame));
    }
  }

  LOG(INFO) << StringPrintf(
      " %d in %.3fs", frames_.size(), timer.ElapsedSeconds());

  //////////////////////////////////////////////////////////////////////////////
  // Load matches
  //////////////////////////////////////////////////////////////////////////////

  timer.Restart();
  LOG(INFO) << "Loading matches...";

  std::vector<std::pair<image_pair_t, TwoViewGeometry>> two_view_geometries =
      database.ReadTwoViewGeometries();

  LOG(INFO) << StringPrintf(
      " %d in %.3fs", two_view_geometries.size(), timer.ElapsedSeconds());

  //////////////////////////////////////////////////////////////////////////////
  // Load images
  //////////////////////////////////////////////////////////////////////////////

  timer.Restart();
  LOG(INFO) << "Loading images...";

  std::unordered_set<frame_t> frame_ids;
  std::unordered_map<image_t, frame_t> image_to_frame_id;

  {
    std::vector<class Image> images = database.ReadAllImages();
    const size_t num_images = images.size();

    for (auto& image : images) {
      // For backwards compatibility with old databases from before having
      // support for rigs/frames, we create a frame for each image.
      if (has_frames) {
        THROW_CHECK(image.HasFrameId());
      } else {
        class Frame frame;
        frame.SetFrameId(image.ImageId());
        frame.SetRigId(image.CameraId());
        frame.AddDataId(image.DataId());
        image.SetFrameId(frame.FrameId());
        frames_.emplace(frame.FrameId(), std::move(frame));
      }

      image_to_frame_id.emplace(image.ImageId(), image.FrameId());
    }

    // Determines for which images data should be loaded.
    if (options.image_names.empty()) {
      for (const auto& image : images) {
        frame_ids.insert(image.FrameId());
      }
    } else {
      for (const auto& image : images) {
        if (options.image_names.count(image.Name()) > 0) {
          frame_ids.insert(image.FrameId());
        }
      }
    }

    // Collect all images that are connected in the correspondence graph.
    std::unordered_set<frame_t> connected_frame_ids;
    if (!options.load_all_images) {
      connected_frame_ids.reserve(frame_ids.size());
      for (const auto& [pair_id, two_view_geometry] : two_view_geometries) {
        if (UseInlierMatchesCheck(options,
                                  two_view_geometry.config,
                                  two_view_geometry.inlier_matches.size())) {
          const auto [image_id1, image_id2] = PairIdToImagePair(pair_id);
          const frame_t frame_id1 = image_to_frame_id.at(image_id1);
          const frame_t frame_id2 = image_to_frame_id.at(image_id2);
          if (frame_ids.count(frame_id1) > 0 &&
              frame_ids.count(frame_id2) > 0) {
            connected_frame_ids.insert(frame_id1);
            connected_frame_ids.insert(frame_id2);
          }
        }
      }
    }

    const std::unordered_set<frame_t>& load_frame_ids =
        options.load_all_images ? frame_ids : connected_frame_ids;

    // Remove frames that should not be loaded.
    for (auto it = frames_.begin(); it != frames_.end();) {
      if (load_frame_ids.count(it->first) == 0) {
        it = frames_.erase(it);
      } else {
        ++it;
      }
    }

    // Load images and their keypoints. When load_all_images is false, only
    // images with correspondences are loaded, as images without matches are
    // not useful for SfM. When load_all_images is true, all candidate images
    // are loaded so that their keypoints are populated (e.g., for
    // triangulation on an existing reconstruction).
    images_.reserve(load_frame_ids.size());
    for (auto& image : images) {
      if (load_frame_ids.count(image.FrameId()) == 0) {
        continue;
      }

      const image_t image_id = image.ImageId();
      image.SetPoints2D(
          FeatureKeypointsToPointsVector(database.ReadKeypoints(image_id)));
      images_.emplace(image_id, std::move(image));
    }

    if (options.load_all_images) {
      LOG(INFO) << StringPrintf(" %d in %.3fs (loaded all %d)",
                                num_images,
                                timer.ElapsedSeconds(),
                                images_.size());
    } else {
      LOG(INFO) << StringPrintf(" %d in %.3fs (connected %d, loaded %d)",
                                num_images,
                                timer.ElapsedSeconds(),
                                connected_frame_ids.size(),
                                images_.size());
    }
  }

  //////////////////////////////////////////////////////////////////////////////
  // Load pose priors
  //////////////////////////////////////////////////////////////////////////////

  timer.Restart();

  LOG(INFO) << "Loading pose priors...";

  pose_priors_ = database.ReadAllPosePriors();

  if (options.convert_pose_priors_to_enu) {
    ConvertPosePriorsToENU();
  }

  LOG(INFO) << StringPrintf(
      " %d in %.3fs", pose_priors_.size(), timer.ElapsedSeconds());

  //////////////////////////////////////////////////////////////////////////////
  // Build correspondence graph
  //////////////////////////////////////////////////////////////////////////////

  timer.Restart();
  LOG(INFO) << "Building correspondence graph...";

  correspondence_graph_ = std::make_shared<class CorrespondenceGraph>();

  for (const auto& [image_id, image] : images_) {
    correspondence_graph_->AddImage(image_id, image.NumPoints2D());
  }

  size_t num_ignored_image_pairs = 0;
  for (auto& [pair_id, two_view_geometry] : two_view_geometries) {
    if (UseInlierMatchesCheck(options,
                              two_view_geometry.config,
                              two_view_geometry.inlier_matches.size())) {
      const auto [image_id1, image_id2] = PairIdToImagePair(pair_id);
      const frame_t frame_id1 = image_to_frame_id.at(image_id1);
      const frame_t frame_id2 = image_to_frame_id.at(image_id2);
      if (frame_ids.count(frame_id1) > 0 && frame_ids.count(frame_id2) > 0) {
        correspondence_graph_->AddTwoViewGeometry(
            image_id1, image_id2, std::move(two_view_geometry));
      } else {
        num_ignored_image_pairs += 1;
      }
    } else {
      num_ignored_image_pairs += 1;
    }
  }

  correspondence_graph_->Finalize();

  LOG(INFO) << StringPrintf(" in %.3fs (ignored %d)",
                            timer.ElapsedSeconds(),
                            num_ignored_image_pairs);
}

std::shared_ptr<DatabaseCache> DatabaseCache::Create(const Database& database,
                                                     const Options& options) {
  auto cache = std::make_shared<DatabaseCache>();
  cache->Load(database, options);
  return cache;
}

std::shared_ptr<DatabaseCache> DatabaseCache::CreateFromCache(
    const DatabaseCache& database_cache, const Options& options) {
  auto cache = std::make_shared<DatabaseCache>();

  // Collect candidate image ids matching the name filter.
  // Empty image_names means use all images.
  std::unordered_set<image_t> candidate_image_ids;
  for (const auto& [image_id, image] : database_cache.Images()) {
    if (options.image_names.empty() ||
        options.image_names.count(image.Name()) > 0) {
      candidate_image_ids.insert(image_id);
    }
  }

  const auto& source_graph = database_cache.CorrespondenceGraph();

  std::unordered_set<image_t> connected_image_ids;
  if (!options.load_all_images) {
    for (const auto& [pair_id, num_matches] :
         source_graph->NumMatchesBetweenAllImages()) {
      const auto [image_id1, image_id2] = PairIdToImagePair(pair_id);
      if (candidate_image_ids.count(image_id1) == 0 ||
          candidate_image_ids.count(image_id2) == 0) {
        continue;
      }
      const TwoViewGeometry two_view_geometry =
          source_graph->ExtractTwoViewGeometry(
              image_id1, image_id2, /*extract_inlier_matches=*/false);
      if (!UseInlierMatchesCheck(
              options, two_view_geometry.config, num_matches)) {
        continue;
      }
      connected_image_ids.insert(image_id1);
      connected_image_ids.insert(image_id2);
    }
  }

  const std::unordered_set<image_t>& load_image_ids =
      options.load_all_images ? candidate_image_ids : connected_image_ids;

  // Collect frame ids for images to load.
  std::unordered_set<frame_t> filtered_frame_ids;
  for (const image_t image_id : load_image_ids) {
    const auto& image = database_cache.Image(image_id);
    filtered_frame_ids.insert(image.FrameId());
  }

  // Copy all images of filtered frames (not just the images matching the
  // name filter). This is needed for multi-camera rigs where the generalized
  // pose solver needs all images of a frame.
  std::unordered_set<camera_t> filtered_camera_ids;
  for (const auto& [image_id, image] : database_cache.Images()) {
    if (filtered_frame_ids.count(image.FrameId()) > 0) {
      cache->images_.emplace(image_id, image);
      filtered_camera_ids.insert(image.CameraId());
    }
  }

  // Copy filtered frames and collect rig ids.
  std::unordered_set<rig_t> filtered_rig_ids;
  for (const auto& [frame_id, frame] : database_cache.Frames()) {
    if (filtered_frame_ids.count(frame_id) > 0) {
      cache->frames_.emplace(frame_id, frame);
      filtered_rig_ids.insert(frame.RigId());
    }
  }

  // Copy filtered cameras.
  for (const auto& [camera_id, camera] : database_cache.Cameras()) {
    if (filtered_camera_ids.count(camera_id) > 0) {
      cache->cameras_.emplace(camera_id, camera);
    }
  }

  // Copy filtered rigs.
  for (const auto& [rig_id, rig] : database_cache.Rigs()) {
    if (filtered_rig_ids.count(rig_id) > 0) {
      cache->rigs_.emplace(rig_id, rig);
    }
  }

  // Copy pose priors.
  cache->pose_priors_ = database_cache.PosePriors();
  if (options.convert_pose_priors_to_enu) {
    cache->ConvertPosePriorsToENU();
  }

  // Build filtered correspondence graph with all images from connected frames.
  cache->correspondence_graph_ = std::make_shared<class CorrespondenceGraph>();

  for (const auto& [image_id, image] : cache->images_) {
    cache->correspondence_graph_->AddImage(image_id, image.NumPoints2D());
  }

  // Copy correspondences between all image pairs in the cache.
  for (const image_pair_t pair_id : source_graph->ImagePairs()) {
    const auto [image_id1, image_id2] = PairIdToImagePair(pair_id);
    if (cache->images_.count(image_id1) > 0 &&
        cache->images_.count(image_id2) > 0) {
      cache->correspondence_graph_->AddTwoViewGeometry(
          image_id1,
          image_id2,
          source_graph->ExtractTwoViewGeometry(
              image_id1, image_id2, /*extract_inlier_matches=*/true));
    }
  }

  cache->correspondence_graph_->Finalize();

  return cache;
}

void DatabaseCache::AddRig(class Rig rig) {
  const rig_t rig_id = rig.RigId();
  THROW_CHECK(!ExistsRig(rig_id));
  rigs_.emplace(rig_id, std::move(rig));
}

void DatabaseCache::AddCamera(struct Camera camera) {
  const camera_t camera_id = camera.camera_id;
  THROW_CHECK(!ExistsCamera(camera_id));
  cameras_.emplace(camera_id, std::move(camera));
}

void DatabaseCache::AddFrame(class Frame frame) {
  const rig_t frame_id = frame.FrameId();
  THROW_CHECK(!ExistsFrame(frame_id));
  frames_.emplace(frame_id, std::move(frame));
}

void DatabaseCache::AddImage(class Image image) {
  const image_t image_id = image.ImageId();
  THROW_CHECK(!ExistsImage(image_id));
  correspondence_graph_->AddImage(image_id, image.NumPoints2D());
  images_.emplace(image_id, std::move(image));
}

void DatabaseCache::AddPosePrior(struct PosePrior pose_prior) {
  pose_priors_.push_back(std::move(pose_prior));
}

const class Image* DatabaseCache::FindImageWithName(
    const std::string& name) const {
  for (const auto& image : images_) {
    if (image.second.Name() == name) {
      return &image.second;
    }
  }
  return nullptr;
}

void DatabaseCache::ConvertPosePriorsToENU() {
  bool prior_is_gps = true;

  std::vector<Eigen::Vector3d> gps_prior_positions;
  std::set<PosePrior::CoordinateSystem> coordinate_systems;
  for (const auto& pose_prior : pose_priors_) {
    coordinate_systems.insert(pose_prior.coordinate_system);
    if (pose_prior.coordinate_system != PosePrior::CoordinateSystem::WGS84) {
      prior_is_gps = false;
    } else {
      gps_prior_positions.push_back(pose_prior.position);
    }
  }

  THROW_CHECK_LE(coordinate_systems.size(), 1)
      << "Inconsistent coordinate systems defined in pose priors";

  // If GPS priors are available, convert them to Cartesian ENU coordinates.
  if (prior_is_gps && !gps_prior_positions.empty()) {
    // GPS reference to be used for EllipsoidToENU conversion.
    const double ref_lat = gps_prior_positions[0][0];
    const double ref_lon = gps_prior_positions[0][1];
    const double ref_alt = gps_prior_positions[0][2];

    const GPSTransform gps_transform(GPSTransform::Ellipsoid::WGS84);
    const std::vector<Eigen::Vector3d> v_xyz_prior =
        gps_transform.EllipsoidToENU(
            gps_prior_positions, ref_lat, ref_lon, ref_alt);

    auto xyz_prior_it = v_xyz_prior.begin();
    for (auto& pose_prior : pose_priors_) {
      pose_prior.position = *xyz_prior_it;
      pose_prior.coordinate_system = PosePrior::CoordinateSystem::CARTESIAN;
      ++xyz_prior_it;
    }
  }
}

}  // namespace colmap
