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

#include "colmap/geometry/essential_matrix.h"
#include "colmap/geometry/gps.h"
#include "colmap/math/math.h"
#include "colmap/math/random.h"
#include "colmap/math/union_find.h"
#include "colmap/sensor/bitmap.h"
#include "colmap/util/eigen_alignment.h"

#include <Eigen/Geometry>

namespace colmap {
namespace {

void AddOutlierMatches(double inlier_ratio,
                       int num_points2D1,
                       int num_points2D2,
                       FeatureMatches* matches) {
  const int num_outliers = matches->size() * (1.0 - inlier_ratio);
  for (int i = 0; i < num_outliers; ++i) {
    matches->emplace_back(
        RandomUniformInteger<point2D_t>(0, num_points2D1 - 1),
        RandomUniformInteger<point2D_t>(0, num_points2D2 - 2));
  }
  std::shuffle(matches->begin(), matches->end(), *PRNG);
}

std::vector<image_pair_t> ExtractExhaustiveImagePairs(
    const Reconstruction& reconstruction) {
  const size_t num_exhaustive_pairs =
      reconstruction.NumRegImages() * (reconstruction.NumRegImages() - 1) / 2;
  std::vector<image_pair_t> image_pairs;
  image_pairs.reserve(num_exhaustive_pairs);
  for (const image_t image_id1 : reconstruction.RegImageIds()) {
    for (const image_t image_id2 : reconstruction.RegImageIds()) {
      if (image_id1 >= image_id2) {
        continue;
      }
      image_pairs.push_back(ImagePairToPairId(image_id1, image_id2));
    }
  }
  THROW_CHECK_EQ(image_pairs.size(), num_exhaustive_pairs);
  return image_pairs;
}

// Fill the configuration, essential matrix, and (for perspective pairs)
// fundamental matrix of a synthetic two-view geometry.
void SetTwoViewGeometryModel(const Camera& camera1,
                             const Camera& camera2,
                             const Rigid3d& cam2_from_cam1,
                             TwoViewGeometry* two_view_geometry) {
  two_view_geometry->E = EssentialMatrixFromPose(cam2_from_cam1);
  if (camera1.IsSpherical() || camera2.IsSpherical()) {
    // Omnidirectional cameras (e.g. EQUIRECTANGULAR) have no pinhole
    // calibration matrix, so the fundamental matrix is undefined; they are
    // always calibrated and use the bearing-based essential-matrix
    // configuration.
    two_view_geometry->config = TwoViewGeometry::CALIBRATED;
    return;
  }
  const bool is_calibrated =
      camera1.has_prior_focal_length && camera2.has_prior_focal_length;
  two_view_geometry->config = is_calibrated ? TwoViewGeometry::CALIBRATED
                                            : TwoViewGeometry::UNCALIBRATED;
  two_view_geometry->F =
      FundamentalFromEssentialMatrix(camera2.CalibrationMatrix(),
                                     *two_view_geometry->E,
                                     camera1.CalibrationMatrix());
}

TwoViewGeometry BuildTwoViewGeometry(bool has_relative_pose,
                                     const Reconstruction& reconstruction,
                                     const image_pair_t pair_id) {
  const auto [image_id1, image_id2] = PairIdToImagePair(pair_id);

  const auto& image1 = reconstruction.Image(image_id1);
  const auto num_points2D1 = image1.NumPoints2D();
  const auto& image2 = reconstruction.Image(image_id2);
  const auto num_points2D2 = image2.NumPoints2D();

  TwoViewGeometry two_view_geometry;
  const Rigid3d cam2_from_cam1 =
      image2.CamFromWorld() * Inverse(image1.CamFromWorld());
  if (has_relative_pose) {
    two_view_geometry.cam2_from_cam1 = cam2_from_cam1;
  }
  SetTwoViewGeometryModel(*image1.CameraPtr(),
                          *image2.CameraPtr(),
                          cam2_from_cam1,
                          &two_view_geometry);

  for (point2D_t point2D_idx1 = 0; point2D_idx1 < num_points2D1;
       ++point2D_idx1) {
    const auto& point2D1 = image1.Point2D(point2D_idx1);
    if (!point2D1.HasPoint3D()) {
      continue;
    }
    for (point2D_t point2D_idx2 = 0; point2D_idx2 < num_points2D2;
         ++point2D_idx2) {
      const auto& point2D2 = image2.Point2D(point2D_idx2);
      if (point2D1.point3D_id == point2D2.point3D_id) {
        two_view_geometry.inlier_matches.emplace_back(point2D_idx1,
                                                      point2D_idx2);
        break;
      }
    }
  }
  return two_view_geometry;
}

void WriteTwoViewGeometryToDatabase(const image_pair_t pair_id,
                                    const TwoViewGeometry& two_view_geometry,
                                    double inlier_match_ratio,
                                    const Reconstruction& reconstruction,
                                    Database* database) {
  const auto [image_id1, image_id2] = PairIdToImagePair(pair_id);
  const auto& image1 = reconstruction.Image(image_id1);
  const auto& image2 = reconstruction.Image(image_id2);

  FeatureMatches matches = two_view_geometry.inlier_matches;
  AddOutlierMatches(
      inlier_match_ratio, image1.NumPoints2D(), image2.NumPoints2D(), &matches);

  if (!database->ExistsMatches(image_id1, image_id2)) {
    database->WriteMatches(image_id1, image_id2, matches);
  }
  if (!database->ExistsTwoViewGeometry(image_id1, image_id2)) {
    database->WriteTwoViewGeometry(image_id1, image_id2, two_view_geometry);
  }
}

void SynthesizeExhaustiveMatches(double inlier_match_ratio,
                                 bool has_relative_pose,
                                 Reconstruction* reconstruction,
                                 Database* database) {
  for (const image_pair_t pair_id :
       ExtractExhaustiveImagePairs(*reconstruction)) {
    const TwoViewGeometry two_view_geometry =
        BuildTwoViewGeometry(has_relative_pose, *reconstruction, pair_id);
    WriteTwoViewGeometryToDatabase(pair_id,
                                   two_view_geometry,
                                   inlier_match_ratio,
                                   *reconstruction,
                                   database);
  }
}

void SynthesizeChainedMatches(double inlier_match_ratio,
                              bool has_relative_pose,
                              Reconstruction* reconstruction,
                              Database* database) {
  std::unordered_map<image_pair_t, TwoViewGeometry> two_view_geometries;
  for (const auto& point3D : reconstruction->Points3D()) {
    std::vector<TrackElement> track_elements = point3D.second.track.Elements();
    std::sort(track_elements.begin(),
              track_elements.end(),
              [](const TrackElement& left, const TrackElement& right) {
                return left.image_id < right.image_id;
              });
    for (size_t i = 1; i < track_elements.size(); ++i) {
      const auto& prev_track_el = track_elements[i - 1];
      const auto& curr_track_el = track_elements[i];
      if (curr_track_el.image_id != prev_track_el.image_id + 1) {
        continue;
      }
      const image_pair_t pair_id =
          ImagePairToPairId(prev_track_el.image_id, curr_track_el.image_id);
      if (ShouldSwapImagePair(prev_track_el.image_id, curr_track_el.image_id)) {
        two_view_geometries[pair_id].inlier_matches.emplace_back(
            curr_track_el.point2D_idx, prev_track_el.point2D_idx);
      } else {
        two_view_geometries[pair_id].inlier_matches.emplace_back(
            prev_track_el.point2D_idx, curr_track_el.point2D_idx);
      }
    }
  }

  for (auto& [pair_id, two_view_geometry] : two_view_geometries) {
    const auto [image_id1, image_id2] = PairIdToImagePair(pair_id);
    const auto& image1 = reconstruction->Image(image_id1);
    const auto& camera1 = *image1.CameraPtr();
    const auto& image2 = reconstruction->Image(image_id2);
    const auto& camera2 = *image2.CameraPtr();
    const Rigid3d cam2_from_cam1 =
        image2.CamFromWorld() * Inverse(image1.CamFromWorld());
    if (has_relative_pose) {
      two_view_geometry.cam2_from_cam1 = cam2_from_cam1;
    }
    SetTwoViewGeometryModel(
        camera1, camera2, cam2_from_cam1, &two_view_geometry);

    WriteTwoViewGeometryToDatabase(pair_id,
                                   two_view_geometry,
                                   inlier_match_ratio,
                                   *reconstruction,
                                   database);
  }
}

bool IsViewGraphConnected(const std::set<image_t>& images,
                          const std::set<image_pair_t>& image_pairs) {
  if (images.empty()) {
    return true;
  }
  if (images.size() == 1) {
    return true;
  }
  if (image_pairs.empty()) {
    return false;
  }

  // Use UnionFind to check connectivity.
  UnionFind<image_t> uf;
  uf.Reserve(images.size());

  for (const image_pair_t pair_id : image_pairs) {
    const auto [image_id1, image_id2] = PairIdToImagePair(pair_id);
    uf.Union(image_id1, image_id2);
  }

  // Check that all images have the same root.
  const image_t root = uf.Find(*images.begin());
  for (const image_t node : images) {
    if (uf.Find(node) != root) {
      return false;
    }
  }
  return true;
}

void SynthesizeSparseMatches(double inlier_match_ratio,
                             bool has_relative_pose,
                             double sparsity,
                             Reconstruction* reconstruction,
                             Database* database) {
  THROW_CHECK_GE(sparsity, 0.0);
  THROW_CHECK_LE(sparsity, 1.0);

  if (sparsity == 0.0) {
    SynthesizeExhaustiveMatches(
        inlier_match_ratio, has_relative_pose, reconstruction, database);
    return;
  }

  if (sparsity == 1.0) {
    return;
  }

  std::vector<image_pair_t> remaining_image_pairs =
      ExtractExhaustiveImagePairs(*reconstruction);

  std::set<image_t> all_image_ids;
  for (const image_t image_id : reconstruction->RegImageIds()) {
    all_image_ids.insert(image_id);
  }

  const size_t num_edges_total = remaining_image_pairs.size();
  const size_t num_edges_to_remove =
      static_cast<size_t>(sparsity * num_edges_total);

  // Try to remove edges randomly while maintaining connectivity.
  std::shuffle(
      remaining_image_pairs.begin(), remaining_image_pairs.end(), *PRNG);
  std::set<image_pair_t> remaining_edges_set(remaining_image_pairs.begin(),
                                             remaining_image_pairs.end());
  size_t edges_removed = 0;
  for (const image_pair_t pair_id : remaining_image_pairs) {
    if (edges_removed >= num_edges_to_remove) {
      break;
    }
    remaining_edges_set.erase(pair_id);
    if (IsViewGraphConnected(all_image_ids, remaining_edges_set)) {
      edges_removed++;
    } else {
      remaining_edges_set.insert(pair_id);
    }
  }

  for (const image_pair_t pair_id : remaining_edges_set) {
    const TwoViewGeometry two_view_geometry =
        BuildTwoViewGeometry(has_relative_pose, *reconstruction, pair_id);
    WriteTwoViewGeometryToDatabase(pair_id,
                                   two_view_geometry,
                                   inlier_match_ratio,
                                   *reconstruction,
                                   database);
  }
}

static const GPSTransform kGPSTransform;
static constexpr double kLat0 = 47.37851943807808;
static constexpr double kLon0 = 8.549099927632087;
static constexpr double kAlt0 = 451.5;

void PosePriorPositionCartesianToWGS84(PosePrior& pose_prior) {
  pose_prior.position = kGPSTransform.ENUToEllipsoid(
      {pose_prior.position}, kLat0, kLon0, kAlt0)[0];
  pose_prior.coordinate_system = PosePrior::CoordinateSystem::WGS84;
}

void PosePriorPositionWGS84ToCartesian(PosePrior& pose_prior) {
  pose_prior.position = kGPSTransform.EllipsoidToENU(
      {pose_prior.position}, kLat0, kLon0, kAlt0)[0];
  pose_prior.coordinate_system = PosePrior::CoordinateSystem::CARTESIAN;
}

}  // namespace

void SynthesizeDataset(const SyntheticDatasetOptions& options,
                       Reconstruction* reconstruction,
                       Database* database) {
  THROW_CHECK_GT(options.num_rigs, 0);
  THROW_CHECK_GT(options.num_cameras_per_rig, 0);
  THROW_CHECK_GT(options.num_frames_per_rig, 0);
  THROW_CHECK_GE(options.num_points3D, 0);
  THROW_CHECK(options.track_length == -1 || options.track_length >= 2);
  if (options.track_length > 0) {
    const int num_images = options.num_rigs * options.num_cameras_per_rig *
                           options.num_frames_per_rig;
    if (options.track_length > num_images) {
      LOG(WARNING) << "track_length (" << options.track_length
                   << ") exceeds number of images (" << num_images
                   << "), skipping observation pruning.";
    }
  }
  THROW_CHECK_NE(options.feature_type, FeatureExtractorType::UNDEFINED);
  THROW_CHECK_GE(options.num_points2D_without_point3D, 0);
  THROW_CHECK_GE(options.sensor_from_rig_translation_stddev, 0.);
  THROW_CHECK_GE(options.sensor_from_rig_rotation_stddev, 0.);
  THROW_CHECK_GE(options.match_sparsity, 0.);
  THROW_CHECK_LE(options.match_sparsity, 1.);
  THROW_CHECK(!options.image_extension.empty());
  THROW_CHECK(options.image_extension[0] == '.');

  if (PRNG == nullptr) {
    SetPRNGSeed();
  }

  // Synthesize 3D points on unit sphere centered at origin.
  std::unordered_set<point3D_t> new_points3D_ids;
  new_points3D_ids.reserve(options.num_points3D);
  for (int point3D_idx = 0; point3D_idx < options.num_points3D; ++point3D_idx) {
    new_points3D_ids.insert(
        reconstruction->AddPoint3D(Eigen::Vector3d::Random().normalized(),
                                   /*track=*/{}));
  }

  int total_num_images = (database == nullptr) ? 0 : database->NumImages();
  int total_num_descriptors =
      (database == nullptr) ? 0 : database->NumDescriptors();

  int desc_dim = 0;
  switch (options.feature_type) {
    case FeatureExtractorType::SIFT:
      desc_dim = 128;
      break;
    case FeatureExtractorType::UNDEFINED:
      desc_dim = 0;
      break;
    default:
      LOG(FATAL_THROW) << "Invalid FeatureExtractorType specified: "
                       << static_cast<int>(options.feature_type);
  }

  for (int rig_idx = 0; rig_idx < options.num_rigs; ++rig_idx) {
    Rig rig;

    std::vector<sensor_t> camera_sensor_ids;
    camera_sensor_ids.reserve(options.num_cameras_per_rig);
    for (int camera_idx = 0; camera_idx < options.num_cameras_per_rig;
         ++camera_idx) {
      Camera camera;
      camera.width = options.camera_width;
      camera.height = options.camera_height;
      camera.model_id = options.camera_model_id;
      camera.params = options.camera_params;
      THROW_CHECK(camera.VerifyParams());
      camera.has_prior_focal_length = options.camera_has_prior_focal_length;
      camera.camera_id =
          (database == nullptr)
              ? (rig_idx * options.num_cameras_per_rig + camera_idx + 1)
              : database->WriteCamera(camera);
      reconstruction->AddCamera(camera);

      if (rig.NumSensors() == 0) {
        rig.AddRefSensor(camera.SensorId());
      } else {
        Rigid3d sensor_from_rig;
        if (options.sensor_from_rig_rotation_stddev > 0) {
          // Generate a random rotation around the Z-axis.
          // This is to avoid 2D points fall behind the camera.
          const double angle =
              std::clamp(RandomGaussian<double>(
                             0, options.sensor_from_rig_rotation_stddev),
                         -180.0,
                         180.0);
          sensor_from_rig.rotation() = Eigen::Quaterniond(
              Eigen::AngleAxisd(DegToRad(angle), Eigen::Vector3d(0, 0, 1)));
        }
        if (options.sensor_from_rig_translation_stddev > 0) {
          sensor_from_rig.translation() = Eigen::Vector3d(
              RandomGaussian<double>(
                  0, options.sensor_from_rig_translation_stddev),
              RandomGaussian<double>(
                  0, options.sensor_from_rig_translation_stddev),
              RandomGaussian<double>(
                  0, options.sensor_from_rig_translation_stddev));
        }
        rig.AddSensor(camera.SensorId(), sensor_from_rig);
      }

      camera_sensor_ids.push_back(camera.SensorId());
    }

    const rig_t rig_id =
        (database == nullptr) ? rig_idx + 1 : database->WriteRig(rig);
    rig.SetRigId(rig_id);
    reconstruction->AddRig(rig);

    for (int frame_idx = 0; frame_idx < options.num_frames_per_rig;
         ++frame_idx) {
      Frame frame;
      frame.SetRigId(rig.RigId());

      // Synthesize frames as sphere centered at world origin.
      const Eigen::Vector3d view_dir = -Eigen::Vector3d::Random().normalized();
      const Eigen::Vector3d proj_center = -5 * view_dir;
      Rigid3d rig_from_world;
      rig_from_world.rotation() = Eigen::Quaterniond::FromTwoVectors(
          view_dir, Eigen::Vector3d(0, 0, 1));
      rig_from_world.translation() = rig_from_world.rotation() * -proj_center;

      frame.SetRigFromWorld(rig_from_world);

      std::vector<Image> images;
      images.reserve(options.num_cameras_per_rig);
      std::vector<Rigid3d> cams_from_world;
      cams_from_world.reserve(options.num_cameras_per_rig);
      for (const auto& sensor_id : camera_sensor_ids) {
        ++total_num_images;

        Image& image = images.emplace_back();
        image.SetName(StringPrintf("camera%06d_frame%06d%s",
                                   sensor_id.id,
                                   frame_idx,
                                   options.image_extension.c_str()));
        image.SetCameraId(sensor_id.id);
        const image_t image_id = (database == nullptr)
                                     ? total_num_images
                                     : database->WriteImage(image);
        image.SetImageId(image_id);

        frame.AddDataId(image.DataId());

        // Need to compose cam_from_world manually, because the frame/rig
        // pointer references are not yet set up.
        const Camera& camera = reconstruction->Camera(image.CameraId());
        const Rigid3d sensor_from_rig =
            rig.IsRefSensor(camera.SensorId())
                ? Rigid3d()
                : rig.SensorFromRig(camera.SensorId());
        const Rigid3d cam_from_world = sensor_from_rig * rig_from_world;
        cams_from_world.push_back(cam_from_world);

        if (options.prior_position || options.prior_gravity) {
          PosePrior pose_prior;

          if (options.prior_position) {
            pose_prior.position = cam_from_world.TgtOriginInSrc();
            pose_prior.coordinate_system =
                PosePrior::CoordinateSystem::CARTESIAN;
            switch (options.prior_position_coordinate_system) {
              case PosePrior::CoordinateSystem::CARTESIAN:
                break;
              case PosePrior::CoordinateSystem::WGS84:
                PosePriorPositionCartesianToWGS84(pose_prior);
                break;
              default:
                LOG(FATAL) << "Invalid PosePrior::CoordinateSystem specified";
            }
          }

          if (options.prior_gravity) {
            pose_prior.gravity =
                (cam_from_world.rotation() * options.prior_gravity_in_world)
                    .normalized();
          }

          pose_prior.corr_data_id = image.DataId();
          pose_prior.pose_prior_id = database->WritePosePrior(pose_prior);
        }
      }

      const frame_t frame_id =
          (database == nullptr)
              ? (rig_idx * options.num_frames_per_rig + frame_idx + 1)
              : database->WriteFrame(frame);
      frame.SetFrameId(frame_id);
      reconstruction->AddFrame(std::move(frame));

      for (int camera_idx = 0; camera_idx < options.num_cameras_per_rig;
           ++camera_idx) {
        Image& image = images[camera_idx];
        const Camera& camera = reconstruction->Camera(image.CameraId());
        const Rigid3d& cam_from_world = cams_from_world[camera_idx];

        image.SetFrameId(frame_id);

        std::vector<Point2D> points2D;
        points2D.reserve(options.num_points3D +
                         options.num_points2D_without_point3D);

        // Create 3D point observations by projecting 3D points to the image.
        for (auto& [point3D_id, point3D] : reconstruction->Points3D()) {
          if (new_points3D_ids.count(point3D_id) == 0) {
            // If a non-empty reconstruction is given, only add tracks for
            // newly added images and 3D points.
            continue;
          }

          Point2D point2D;
          const std::optional<Eigen::Vector2d> proj_point2D =
              camera.ImgFromCam(cam_from_world * point3D.xyz);
          THROW_CHECK(proj_point2D.has_value());
          point2D.xy = *proj_point2D;
          if (point2D.xy(0) >= 0 && point2D.xy(1) >= 0 &&
              point2D.xy(0) <= camera.width && point2D.xy(1) <= camera.height) {
            point2D.point3D_id = point3D_id;
            points2D.push_back(point2D);
          }
        }

        // Synthesize uniform random 2D points without 3D points.
        for (int i = 0; i < options.num_points2D_without_point3D; ++i) {
          Point2D point2D;
          point2D.xy =
              Eigen::Vector2d(RandomUniformReal<double>(0, camera.width),
                              RandomUniformReal<double>(0, camera.height));
          points2D.push_back(point2D);
        }

        // Shuffle 2D points, so each image has 3D points ordered differently.
        std::shuffle(points2D.begin(), points2D.end(), *PRNG);

        if (database != nullptr) {
          // Create keypoints to add to database.
          FeatureKeypoints keypoints;
          keypoints.reserve(points2D.size());
          FeatureDescriptors descriptors;
          descriptors.type = options.feature_type;
          descriptors.data.resize(points2D.size(), desc_dim);
          std::uniform_int_distribution<int> feature_distribution(0, 255);
          for (point2D_t point2D_idx = 0; point2D_idx < points2D.size();
               ++point2D_idx) {
            const auto& point2D = points2D[point2D_idx];
            keypoints.emplace_back(point2D.xy(0), point2D.xy(1));
            // Generate a unique descriptor for each 3D point. If the 2D point
            // does not observe a 3D point, generate a random unique
            // descriptor.
            std::mt19937 feature_generator(point2D.HasPoint3D()
                                               ? point2D.point3D_id
                                               : options.num_points3D +
                                                     (++total_num_descriptors));
            for (int d = 0; d < descriptors.data.cols(); ++d) {
              descriptors.data(point2D_idx, d) =
                  feature_distribution(feature_generator);
            }
          }
          database->WriteKeypoints(image.ImageId(), keypoints);
          database->WriteDescriptors(image.ImageId(), descriptors);
        }

        for (point2D_t point2D_idx = 0; point2D_idx < points2D.size();
             ++point2D_idx) {
          const auto& point2D = points2D[point2D_idx];
          if (point2D.HasPoint3D()) {
            auto& point3D = reconstruction->Point3D(point2D.point3D_id);
            point3D.track.AddElement(image.ImageId(), point2D_idx);
          }
        }

        image.SetPoints2D(std::move(points2D));

        if (database != nullptr) {
          database->UpdateImage(image);
        }
        reconstruction->AddImage(image);
      }
    }
  }

  if (database != nullptr) {
    switch (options.match_config) {
      case SyntheticDatasetOptions::MatchConfig::EXHAUSTIVE:
        SynthesizeExhaustiveMatches(options.inlier_match_ratio,
                                    options.two_view_geometry_has_relative_pose,
                                    reconstruction,
                                    database);
        break;
      case SyntheticDatasetOptions::MatchConfig::CHAINED:
        SynthesizeChainedMatches(options.inlier_match_ratio,
                                 options.two_view_geometry_has_relative_pose,
                                 reconstruction,
                                 database);
        break;
      case SyntheticDatasetOptions::MatchConfig::SPARSE:
        SynthesizeSparseMatches(options.inlier_match_ratio,
                                options.two_view_geometry_has_relative_pose,
                                options.match_sparsity,
                                reconstruction,
                                database);
        break;
      default:
        LOG(FATAL_THROW) << "Invalid MatchConfig specified";
    }
  }

  if (options.track_length > 0) {
    std::vector<point3D_t> point3D_ids;
    point3D_ids.reserve(reconstruction->NumPoints3D());
    for (const auto& [point3D_id, _] : reconstruction->Points3D()) {
      point3D_ids.push_back(point3D_id);
    }
    for (const point3D_t point3D_id : point3D_ids) {
      const auto& track = reconstruction->Point3D(point3D_id).track;
      if (static_cast<int>(track.Length()) <= options.track_length) {
        continue;
      }
      auto elements = track.Elements();
      std::shuffle(elements.begin(), elements.end(), *PRNG);
      const int num_to_delete =
          static_cast<int>(elements.size()) - options.track_length;
      for (int i = 0; i < num_to_delete; ++i) {
        reconstruction->DeleteObservation(elements[i].image_id,
                                          elements[i].point2D_idx);
      }
    }
  }

  reconstruction->UpdatePoint3DErrors();
}

void SynthesizeNoise(const SyntheticNoiseOptions& options,
                     Reconstruction* reconstruction,
                     Database* database) {
  THROW_CHECK_GE(options.rig_from_world_translation_stddev, 0.);
  THROW_CHECK_GE(options.rig_from_world_rotation_stddev, 0.);
  THROW_CHECK_GE(options.point3D_stddev, 0.);
  THROW_CHECK_GE(options.point2D_stddev, 0.);
  THROW_CHECK_GE(options.prior_position_stddev, 0.);
  THROW_CHECK_GE(options.prior_gravity_stddev, 0.);

  for (const frame_t frame_id : reconstruction->RegFrameIds()) {
    Rigid3d& rig_from_world = reconstruction->Frame(frame_id).RigFromWorld();

    if (options.rig_from_world_rotation_stddev > 0.0) {
      const double angle = std::clamp(
          RandomGaussian<double>(0, options.rig_from_world_rotation_stddev),
          -180.0,
          180.0);
      rig_from_world.rotation() *= Eigen::Quaterniond(
          Eigen::AngleAxisd(DegToRad(angle), Eigen::Vector3d::UnitZ()));
    }

    if (options.rig_from_world_translation_stddev > 0.0) {
      rig_from_world.translation() += Eigen::Vector3d(
          RandomGaussian<double>(0, options.rig_from_world_translation_stddev),
          RandomGaussian<double>(0, options.rig_from_world_translation_stddev),
          RandomGaussian<double>(0, options.rig_from_world_translation_stddev));
    }
  }

  if (options.point2D_stddev > 0.0) {
    for (const auto& [image_id, _] : reconstruction->Images()) {
      Image& image = reconstruction->Image(image_id);
      for (auto& point2D : image.Points2D()) {
        point2D.xy +=
            Eigen::Vector2d(RandomGaussian<double>(0, options.point2D_stddev),
                            RandomGaussian<double>(0, options.point2D_stddev));
      }
      if (database != nullptr) {
        std::vector<FeatureKeypoint> keypoints =
            database->ReadKeypoints(image_id);
        for (point2D_t point2D_idx = 0; point2D_idx < keypoints.size();
             ++point2D_idx) {
          keypoints[point2D_idx].x = image.Point2D(point2D_idx).xy(0);
          keypoints[point2D_idx].y = image.Point2D(point2D_idx).xy(1);
        }
        database->UpdateKeypoints(image.ImageId(), keypoints);
      }
    }
  }

  if (options.point3D_stddev > 0.0) {
    for (auto& [point3D_id, _] : reconstruction->Points3D()) {
      reconstruction->Point3D(point3D_id).xyz +=
          Eigen::Vector3d(RandomGaussian<double>(0, options.point3D_stddev),
                          RandomGaussian<double>(0, options.point3D_stddev),
                          RandomGaussian<double>(0, options.point3D_stddev));
    }
  }

  if (database != nullptr && (options.prior_position_stddev > 0.0 ||
                              options.prior_gravity_stddev > 0.0)) {
    for (auto& pose_prior : database->ReadAllPosePriors()) {
      if (options.prior_position_stddev > 0.) {
        const bool prior_in_wgs84 =
            pose_prior.coordinate_system == PosePrior::CoordinateSystem::WGS84;
        if (prior_in_wgs84) {
          PosePriorPositionWGS84ToCartesian(pose_prior);
        }
        pose_prior.position += Eigen::Vector3d(
            RandomGaussian<double>(0, options.prior_position_stddev),
            RandomGaussian<double>(0, options.prior_position_stddev),
            RandomGaussian<double>(0, options.prior_position_stddev));
        if (!pose_prior.HasPositionCov()) {
          pose_prior.position_covariance = Eigen::Matrix3d::Zero();
        }
        pose_prior.position_covariance += options.prior_position_stddev *
                                          options.prior_position_stddev *
                                          Eigen::Matrix3d::Identity();
        if (prior_in_wgs84) {
          PosePriorPositionCartesianToWGS84(pose_prior);
        }
      }
      if (options.prior_gravity_stddev > 0.) {
        const double angle =
            RandomGaussian<double>(0, DegToRad(options.prior_gravity_stddev));
        const Eigen::Vector3d axis =
            pose_prior.gravity.cross(Eigen::Vector3d::Random()).normalized();
        pose_prior.gravity =
            (Eigen::AngleAxisd(angle, axis) * pose_prior.gravity).normalized();
      }
      database->UpdatePosePrior(pose_prior);
    }
  }

  reconstruction->UpdatePoint3DErrors();
}

void SynthesizeImages(const SyntheticImageOptions& options,
                      const Reconstruction& reconstruction,
                      const std::filesystem::path& image_path) {
  THROW_CHECK_GT(options.feature_patch_radius, 0);
  THROW_CHECK_LT(options.feature_peak_radius, options.feature_patch_radius);
  THROW_CHECK_GT(options.feature_patch_max_brightness, 0);
  THROW_CHECK_LT(options.feature_patch_max_brightness, 255);

  const double patch_radius = std::sqrt(2 * options.feature_patch_radius *
                                        options.feature_patch_radius);

  int total_num_descriptors = 0;
  for (const auto& [image_id, image] : reconstruction.Images()) {
    const Camera& camera = *image.CameraPtr();

    Bitmap bitmap(camera.width, camera.height, /*as_rgb=*/true);
    bitmap.Fill(BitmapColor<uint8_t>(0, 0, 0));

    for (const auto& point2D : image.Points2D()) {
      const int x = static_cast<int>(std::round(point2D.xy(0)));
      const int y = static_cast<int>(std::round(point2D.xy(1)));
      if (x < 0 || y < 0 || x >= static_cast<int>(camera.width) ||
          y >= static_cast<int>(camera.height)) {
        continue;
      }

      std::mt19937 feature_generator(point2D.HasPoint3D()
                                         ? point2D.point3D_id
                                         : reconstruction.NumPoints3D() +
                                               (++total_num_descriptors));

      // Draw a circular patch around the feature with a unique pattern with
      // the aim of producing a unique feature descriptor. Make the pattern a
      // bit darker than the peak, so the keypoint is detected at the center.
      const int patch_minx = std::max(x - options.feature_patch_radius, 0);
      const int patch_maxx = std::min(x + options.feature_patch_radius,
                                      static_cast<int>(camera.width));
      const int patch_miny = std::max(y - options.feature_patch_radius, 0);
      const int patch_maxy = std::min(y + options.feature_patch_radius,
                                      static_cast<int>(camera.height));
      for (int py = patch_miny; py < patch_maxy; ++py) {
        for (int px = patch_minx; px < patch_maxx; ++px) {
          const double radius =
              std::sqrt((px - x) * (px - x) + (py - y) * (py - y));
          if (radius > options.feature_patch_radius) {
            continue;
          }
          // Adjust the brightness so it fades out to the edge of the patch.
          std::uniform_int_distribution<int> patch_brightness_distribution(
              0,
              (1.0 - radius / patch_radius) *
                  options.feature_patch_max_brightness);
          bitmap.SetPixel(px,
                          py,
                          BitmapColor<uint8_t>(patch_brightness_distribution(
                              feature_generator)));
        }
      }

      // Draw a small, bright peak around the feature for keypoint detection.
      const int peak_minx = std::max(x - options.feature_peak_radius, 0);
      const int peak_maxx = std::min(x + options.feature_peak_radius,
                                     static_cast<int>(camera.width));
      const int peak_miny = std::max(y - options.feature_peak_radius, 0);
      const int peak_maxy = std::min(y + options.feature_peak_radius,
                                     static_cast<int>(camera.height));
      std::uniform_int_distribution<int> peak_color_distribution(
          options.feature_patch_max_brightness, 255);
      const BitmapColor<uint8_t> peak_color(
          peak_color_distribution(feature_generator),
          peak_color_distribution(feature_generator),
          peak_color_distribution(feature_generator));
      for (int py = peak_miny; py < peak_maxy; ++py) {
        for (int px = peak_minx; px < peak_maxx; ++px) {
          bitmap.SetPixel(px, py, peak_color);
        }
      }
    }

    const auto output_image_path = image_path / image.Name();
    if (!bitmap.Write(output_image_path)) {
      LOG(ERROR) << "Failed to write image to " << output_image_path;
    }
  }
}

}  // namespace colmap
