#include "colmap/controllers/image_reader.h"
#include "colmap/controllers/undistorters.h"
#include "colmap/exe/feature.h"
#include "colmap/image/undistortion.h"
#include "colmap/scene/camera.h"
#include "colmap/scene/reconstruction.h"
#include "colmap/util/base_controller.h"
#include "colmap/util/file.h"
#include "colmap/util/logging.h"

#include "pycolmap/helpers.h"
#include "pycolmap/pybind11_extension.h"

#include <filesystem>
#include <memory>

#include <glog/logging.h>
#include <pybind11/pybind11.h>
#include <pybind11/stl.h>

using namespace colmap;
using namespace pybind11::literals;
namespace py = pybind11;

void ImportImages(const std::filesystem::path& database_path,
                  const std::filesystem::path& image_path,
                  const CameraMode camera_mode,
                  const std::vector<std::string>& image_names,
                  const ImageReaderOptions& options_) {
  THROW_CHECK_FILE_EXISTS(database_path);
  THROW_CHECK_DIR_EXISTS(image_path);

  ImageReaderOptions options(options_);
  options.image_path = image_path;
  options.image_names = image_names;
  UpdateImageReaderOptionsFromCameraMode(options, camera_mode);

  auto database = Database::Open(database_path);
  ImageReader image_reader(options, database.get());

  PyInterrupt py_interrupt(2.0);

  while (image_reader.NextIndex() < image_reader.NumImages()) {
    if (py_interrupt.Raised()) {
      throw py::error_already_set();
    }
    Rig rig;
    Camera camera;
    Image image;
    PosePrior pose_prior;
    Bitmap bitmap;
    const ImageReader::Status status =
        image_reader.Next(&rig, &camera, &image, &pose_prior, &bitmap, nullptr);
    if (status != ImageReader::Status::SUCCESS) {
      LOG(ERROR) << image.Name() << " " << ImageReader::StatusToString(status);
      continue;
    }
    DatabaseTransaction database_transaction(database.get());
    if (image.ImageId() == kInvalidImageId) {
      image.SetImageId(database->WriteImage(image));

      Frame frame;
      frame.SetRigId(rig.RigId());
      frame.AddDataId(image.DataId());

      if (pose_prior.HasPosition() || pose_prior.HasGravity()) {
        pose_prior.corr_data_id = image.DataId();
        database->WritePosePrior(pose_prior);
      }

      database->WriteFrame(frame);
    }
  }
}

Camera InferCameraFromImage(const std::filesystem::path& image_path,
                            const ImageReaderOptions& options) {
  Bitmap bitmap;
  THROW_CHECK_FILE_EXISTS(image_path);
  THROW_CHECK(bitmap.Read(image_path, false))
      << "Cannot read image file: " << image_path;

  const std::optional<double> maybe_focal_length = bitmap.ExifFocalLength();
  const double focal_length =
      maybe_focal_length.value_or(options.default_focal_length_factor *
                                  std::max(bitmap.Width(), bitmap.Height()));
  Camera camera = Camera::CreateFromModelName(kInvalidCameraId,
                                              options.camera_model,
                                              focal_length,
                                              bitmap.Width(),
                                              bitmap.Height());
  camera.has_prior_focal_length = maybe_focal_length.has_value();
  THROW_CHECK(camera.VerifyParams())
      << "Invalid camera params: " << camera.ParamsToString();

  return camera;
}

void UndistortImages(const std::filesystem::path& output_path,
                     const std::filesystem::path& input_path,
                     const std::filesystem::path& image_path,
                     const std::vector<std::string>& image_names,
                     const std::string& output_type,
                     const FileCopyType copy_type,
                     const int num_patch_match_src_images,
                     const UndistortCameraOptions& undistort_camera_options,
                     int jpeg_quality,
                     int num_threads) {
  THROW_CHECK_DIR_EXISTS(image_path);
  CreateDirIfNotExists(output_path);

  if (output_type != "COLMAP") {
    LOG_IF(WARNING, !image_names.empty())
        << "Ignoring `image_names` for output type " << output_type;
    LOG_IF(WARNING, num_patch_match_src_images != -1)
        << "Ignoring `num_patch_match_src_images` for output type "
        << output_type;
  }

  Reconstruction reconstruction;
  reconstruction.Read(input_path);
  LOG(INFO) << StringPrintf(" => Reconstruction with %d images and %d points",
                            reconstruction.NumImages(),
                            reconstruction.NumPoints3D());

  std::unique_ptr<BaseController> undistorter;
  if (output_type == "COLMAP") {
    COLMAPUndistorter::Options options;
    options.copy_type = copy_type;
    if (num_patch_match_src_images > 0) {
      options.num_patch_match_src_images = num_patch_match_src_images;
    }
    options.jpeg_quality = jpeg_quality;
    options.num_threads = num_threads;
    options.image_ids.reserve(image_names.size());
    for (const auto& image_name : image_names) {
      const Image* image = reconstruction.FindImageWithName(image_name);
      if (image != nullptr) {
        options.image_ids.push_back(image->ImageId());
      } else {
        LOG(WARNING) << "Cannot find image " << image_name;
      }
    }
    undistorter = std::make_unique<COLMAPUndistorter>(std::move(options),
                                                      undistort_camera_options,
                                                      reconstruction,
                                                      image_path,
                                                      output_path);
  } else if (output_type == "PMVS") {
    PMVSUndistorter::Options pmvs_options;
    pmvs_options.jpeg_quality = jpeg_quality;
    pmvs_options.num_threads = num_threads;
    undistorter = std::make_unique<PMVSUndistorter>(pmvs_options,
                                                    undistort_camera_options,
                                                    reconstruction,
                                                    image_path,
                                                    output_path);
  } else if (output_type == "CMP-MVS") {
    CMPMVSUndistorter::Options cmpmvs_options;
    cmpmvs_options.jpeg_quality = jpeg_quality;
    cmpmvs_options.num_threads = num_threads;
    undistorter = std::make_unique<CMPMVSUndistorter>(cmpmvs_options,
                                                      undistort_camera_options,
                                                      reconstruction,
                                                      image_path,
                                                      output_path);
  } else {
    LOG(FATAL_THROW)
        << "Invalid `output_type` - supported values are {'COLMAP', "
           "'PMVS', 'CMP-MVS'}.";
  }

  py::gil_scoped_release release;
  undistorter->Run();
}

void BindImages(py::module& m) {
  auto PyCameraMode = py::enum_<CameraMode>(m, "CameraMode")
                          .value("AUTO", CameraMode::AUTO)
                          .value("SINGLE", CameraMode::SINGLE)
                          .value("PER_FOLDER", CameraMode::PER_FOLDER)
                          .value("PER_IMAGE", CameraMode::PER_IMAGE);
  AddStringToEnumConstructor(PyCameraMode);

  using IROpts = ImageReaderOptions;
  auto PyImageReaderOptions =
      py::classh<IROpts>(m, "ImageReaderOptions")
          .def(py::init<>())
          .def_readwrite("camera_model",
                         &IROpts::camera_model,
                         "Name of the camera model.")
          .def_readwrite("mask_path",
                         &IROpts::mask_path,
                         "Optional root path to folder which contains image"
                         "masks. For a given image, the corresponding mask"
                         "must have the same sub-path below this root as the"
                         "image has below image_path. The filename must be"
                         "equal, aside from the added extension .png. "
                         "For example, for an image image_path/abc/012.jpg,"
                         "the mask would be mask_path/abc/012.jpg.png. No"
                         "features will be extracted in regions where the"
                         "mask image is black (pixel intensity value 0 in"
                         "grayscale).")
          .def_readwrite("existing_camera_id",
                         &IROpts::existing_camera_id,
                         "Whether to explicitly use an existing camera for "
                         "all images. Note that in this case the specified "
                         "camera model and parameters are ignored.")
          .def_readwrite("camera_params",
                         &IROpts::camera_params,
                         "Manual specification of camera parameters. If "
                         "empty, camera parameters will be extracted from "
                         "EXIF, i.e. principal point and focal length.")
          .def_readwrite(
              "default_focal_length_factor",
              &IROpts::default_focal_length_factor,
              "If camera parameters are not specified manually and the image "
              "does not have focal length EXIF information, the focal length "
              "is set to the value `default_focal_length_factor * max(width, "
              "height)`.")
          .def_readwrite(
              "camera_mask_path",
              &IROpts::camera_mask_path,
              "Optional path to an image file specifying a mask for all "
              "images. No features will be extracted in regions where the "
              "mask is black (pixel intensity value 0 in grayscale)")
          .def("check", &IROpts::Check);
  MakeDataclass(PyImageReaderOptions);

  auto PyFileCopyType = py::enum_<FileCopyType>(m, "FileCopyType")
                            .value("copy", FileCopyType::COPY)
                            .value("softlink", FileCopyType::SOFT_LINK)
                            .value("hardlink", FileCopyType::HARD_LINK);
  AddStringToEnumConstructor(PyFileCopyType);

  m.def("import_images",
        &ImportImages,
        "database_path"_a,
        "image_path"_a,
        "camera_mode"_a = CameraMode::AUTO,
        "image_names"_a = std::vector<std::string>(),
        py::arg_v("options", ImageReaderOptions(), "ImageReaderOptions()"),
        "Import images into a database");

  m.def("infer_camera_from_image",
        &InferCameraFromImage,
        "image_path"_a,
        py::arg_v("options", ImageReaderOptions(), "ImageReaderOptions()"),
        "Guess the camera parameters from the EXIF metadata");

  m.def("undistort_images",
        &UndistortImages,
        "output_path"_a,
        "input_path"_a,
        "image_path"_a,
        "image_names"_a = std::vector<std::string>(),
        "output_type"_a = "COLMAP",
        "copy_policy"_a = FileCopyType::COPY,
        "num_patch_match_src_images"_a = -1,
        py::arg_v("undistort_options",
                  UndistortCameraOptions(),
                  "UndistortCameraOptions()"),
        "jpeg_quality"_a = -1,
        "num_threads"_a = -1,
        "Undistort images");
}
