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

#include "colmap/estimators/bundle_adjustment_ceres.h"
#include "colmap/estimators/cost_functions/manifold.h"
#include "colmap/estimators/cost_functions/pose_prior.h"
#include "colmap/estimators/cost_functions/reprojection_error.h"
#include "colmap/estimators/cost_functions/utils.h"
#include "colmap/estimators/pose.h"
#include "colmap/estimators/solvers/generalized_absolute_pose.h"
#include "colmap/estimators/solvers/generalized_relative_pose.h"
#include "colmap/geometry/rigid3.h"
#include "colmap/optim/loransac.h"
#include "colmap/optim/support_measurement.h"
#include "colmap/scene/camera.h"
#include "colmap/util/eigen_alignment.h"
#include "colmap/util/logging.h"

#include <Eigen/Core>

namespace colmap {
namespace {

void ThrowCheckCameras(const std::vector<size_t>& camera_idxs,
                       const std::vector<Rigid3d>& cams_from_rig,
                       const std::vector<Camera>& cameras) {
  THROW_CHECK(!cameras.empty());
  THROW_CHECK_EQ(cams_from_rig.size(), cameras.size());
  const auto [min_camera_idx, max_camera_idx] =
      std::minmax_element(camera_idxs.begin(), camera_idxs.end());
  THROW_CHECK_GE(*min_camera_idx, 0);
  THROW_CHECK_LT(*max_camera_idx, cameras.size());
}

bool IsPanoramicRig(const std::vector<size_t>& camera_idxs,
                    const std::vector<Rigid3d>& cams_from_rig) {
  const std::set<size_t> camera_idx_set(camera_idxs.begin(), camera_idxs.end());
  const size_t first_camera_idx = *camera_idx_set.begin();
  const Eigen::Vector3d first_origin_in_rig =
      cams_from_rig[first_camera_idx].TgtOriginInSrc();
  for (auto it = ++camera_idx_set.begin(); it != camera_idx_set.end(); ++it) {
    const Eigen::Vector3d other_origin_in_rig =
        cams_from_rig[*it].TgtOriginInSrc();
    if (!first_origin_in_rig.isApprox(other_origin_in_rig, 1e-6)) {
      return false;
    }
  }
  return true;
}

double ComputeMaxErrorInCamera(const std::vector<size_t>& camera_idxs,
                               const std::vector<Camera>& cameras,
                               const double max_error_px) {
  THROW_CHECK_GT(max_error_px, 0.0);
  double max_error_cam = 0.;
  for (const auto& camera_idx : camera_idxs) {
    max_error_cam += cameras[camera_idx].CamFromImgThreshold(max_error_px);
  }
  return max_error_cam / camera_idxs.size();
}

bool LowerVector3d(const Eigen::Vector3d& v1, const Eigen::Vector3d& v2) {
  if (v1.x() < v2.x()) {
    return true;
  } else if (v1.x() == v2.x()) {
    if (v1.y() < v2.y()) {
      return true;
    } else if (v1.y() == v2.y()) {
      return v1.z() < v2.z();
    } else {
      return false;
    }
  } else {
    return false;
  }
}

std::vector<size_t> ComputeUniquePointIds(
    const std::vector<Eigen::Vector3d>& points3D) {
  std::vector<size_t> point3D_ids(points3D.size());
  std::iota(point3D_ids.begin(), point3D_ids.end(), 0);
  std::sort(point3D_ids.begin(), point3D_ids.end(), [&](size_t i, size_t j) {
    return LowerVector3d(points3D[i], points3D[j]);
  });

  std::vector<size_t>::iterator unique_it = point3D_ids.begin();
  std::vector<size_t>::iterator current_it = point3D_ids.begin();
  std::vector<size_t> unique_point3D_ids(points3D.size());
  while (current_it != point3D_ids.end()) {
    if (!points3D[*unique_it].isApprox(points3D[*current_it], 1e-5)) {
      unique_it = current_it;
    }
    unique_point3D_ids[*current_it] = unique_it - point3D_ids.begin();
    current_it++;
  }
  return unique_point3D_ids;
}

}  // namespace

bool EstimateGeneralizedAbsolutePose(
    const RANSACOptions& options,
    const std::vector<Eigen::Vector2d>& points2D,
    const std::vector<Eigen::Vector3d>& points3D,
    const std::vector<size_t>& camera_idxs,
    const std::vector<Rigid3d>& cams_from_rig,
    const std::vector<Camera>& cameras,
    Rigid3d* rig_from_world,
    size_t* num_inliers,
    std::vector<char>* inlier_mask) {
  THROW_CHECK_EQ(points2D.size(), points3D.size());
  THROW_CHECK_EQ(points2D.size(), camera_idxs.size());
  ThrowCheckCameras(camera_idxs, cams_from_rig, cameras);
  options.Check();
  if (points2D.size() == 0) {
    return false;
  }

  // Precompute cam_from_rig matrices for fast residual computation
  std::vector<Eigen::Matrix3x4d> cams_from_rig_matrices(cams_from_rig.size());
  for (size_t i = 0; i < cams_from_rig.size(); i++) {
    cams_from_rig_matrices[i] = cams_from_rig[i].ToMatrix();
  }

  std::vector<GP3PEstimator::X_t> rig_points2D(points2D.size());
  for (size_t i = 0; i < points2D.size(); i++) {
    const size_t camera_idx = camera_idxs[i];
    if (const std::optional<Eigen::Vector2d> cam_point =
            cameras[camera_idx].CamFromImg(points2D[i]);
        cam_point) {
      rig_points2D[i].ray_in_cam = cam_point->homogeneous().normalized();
    } else {
      rig_points2D[i].ray_in_cam.setZero();
    }
    rig_points2D[i].cam_from_rig = cams_from_rig_matrices[camera_idx];
  }

  // Associate unique ids to each 3D point.
  // Needed for UniqueInlierSupportMeasurer to avoid counting the same
  // 3D point multiple times due to FoV overlap in rig.
  // TODO(sarlinpe): Allow passing unique_point3D_ids as argument.
  std::vector<size_t> unique_point3D_ids = ComputeUniquePointIds(points3D);

  // Average of the errors over the cameras, weighted by the number of
  // correspondences
  RANSACOptions options_copy(options);
  options_copy.max_error =
      ComputeMaxErrorInCamera(camera_idxs, cameras, options.max_error);

  RANSAC<GP3PEstimator, UniqueInlierSupportMeasurer> ransac(
      options_copy,
      GP3PEstimator(GP3PEstimator::ResidualType::ReprojectionError),
      UniqueInlierSupportMeasurer(std::move(unique_point3D_ids)));
  auto report = ransac.Estimate(rig_points2D, points3D);
  if (!report.success) {
    return false;
  }

  *rig_from_world = report.model;
  *num_inliers = report.support.num_unique_inliers;
  *inlier_mask = std::move(report.inlier_mask);

  return true;
}

bool EstimateGeneralizedRelativePose(
    const RANSACOptions& ransac_options,
    const std::vector<Eigen::Vector2d>& points2D1,
    const std::vector<Eigen::Vector2d>& points2D2,
    const std::vector<size_t>& camera_idxs1,
    const std::vector<size_t>& camera_idxs2,
    const std::vector<Rigid3d>& cams_from_rig,
    const std::vector<Camera>& cameras,
    std::optional<Rigid3d>* rig2_from_rig1,
    std::optional<Rigid3d>* pano2_from_pano1,
    size_t* num_inliers,
    std::vector<char>* inlier_mask) {
  THROW_CHECK_EQ(points2D1.size(), points2D2.size());
  ThrowCheckCameras(camera_idxs1, cams_from_rig, cameras);
  ThrowCheckCameras(camera_idxs2, cams_from_rig, cameras);
  ransac_options.Check();

  const size_t num_points = points2D1.size();
  if (num_points == 0) {
    return false;
  }

  // Adjust the error to be in normalized camera coordinates.
  RANSACOptions normalized_ransac_options = ransac_options;
  double normalized_max_error = 0;
  for (const Camera& camera : cameras) {
    normalized_max_error +=
        camera.CamFromImgThreshold(ransac_options.max_error);
  }
  normalized_ransac_options.max_error = normalized_max_error / cameras.size();

  if (IsPanoramicRig(camera_idxs1, cams_from_rig) &&
      IsPanoramicRig(camera_idxs2, cams_from_rig)) {
    Rigid3d cam2_from_cam1;
    std::vector<Eigen::Vector3d> cam_rays1(num_points);
    std::vector<Eigen::Vector3d> cam_rays2(num_points);
    for (size_t i = 0; i < num_points; ++i) {
      const size_t camera_idx1 = camera_idxs1[i];
      if (const std::optional<Eigen::Vector2d> cam_point1 =
              cameras[camera_idx1].CamFromImg(points2D1[i]);
          cam_point1.has_value()) {
        cam_rays1[i] = cams_from_rig[camera_idx1].rotation().inverse() *
                       cam_point1->homogeneous().normalized();
      } else {
        cam_rays1[i].setZero();
      }

      const size_t camera_idx2 = camera_idxs2[i];
      if (const std::optional<Eigen::Vector2d> cam_point2 =
              cameras[camera_idxs2[i]].CamFromImg(points2D2[i]);
          cam_point2.has_value()) {
        cam_rays2[i] = cams_from_rig[camera_idx2].rotation().inverse() *
                       cam_point2->homogeneous().normalized();
      } else {
        cam_rays2[i].setZero();
      }
    }
    if (EstimateRelativePose(normalized_ransac_options,
                             cam_rays1,
                             cam_rays2,
                             &cam2_from_cam1,
                             num_inliers,
                             inlier_mask)) {
      *pano2_from_pano1 = cam2_from_cam1;
      return true;
    }
    return false;
  }

  std::vector<GRNPObservation> points1(num_points);
  std::vector<GRNPObservation> points2(num_points);
  for (size_t i = 0; i < num_points; ++i) {
    points1[i].cam_from_rig = cams_from_rig[camera_idxs1[i]];
    if (const std::optional<Eigen::Vector2d> cam_point1 =
            cameras[camera_idxs1[i]].CamFromImg(points2D1[i]);
        cam_point1.has_value()) {
      points1[i].ray_in_cam = cam_point1->homogeneous().normalized();
    } else {
      points1[i].ray_in_cam.setZero();
    }

    points2[i].cam_from_rig = cams_from_rig[camera_idxs2[i]];
    if (const std::optional<Eigen::Vector2d> cam_point2 =
            cameras[camera_idxs2[i]].CamFromImg(points2D2[i]);
        cam_point2.has_value()) {
      points2[i].ray_in_cam = cam_point2->homogeneous().normalized();
    } else {
      points2[i].ray_in_cam.setZero();
    }
  }

  LORANSAC<GR6PEstimator, GR8PEstimator> ransac(normalized_ransac_options);
  auto report = ransac.Estimate(points1, points2);
  if (!report.success) {
    return false;
  }

  *rig2_from_rig1 = report.model;
  *num_inliers = report.support.num_inliers;
  *inlier_mask = std::move(report.inlier_mask);

  return true;
}

bool RefineGeneralizedAbsolutePose(const AbsolutePoseRefinementOptions& options,
                                   const std::vector<char>& inlier_mask,
                                   const std::vector<Eigen::Vector2d>& points2D,
                                   const std::vector<Eigen::Vector3d>& points3D,
                                   const std::vector<size_t>& camera_idxs,
                                   const std::vector<Rigid3d>& cams_from_rig,
                                   Rigid3d* rig_from_world,
                                   std::vector<Camera>* cameras,
                                   Eigen::Matrix6d* rig_from_world_cov) {
  THROW_CHECK_EQ(points2D.size(), inlier_mask.size());
  THROW_CHECK_EQ(points2D.size(), points3D.size());
  THROW_CHECK_EQ(points2D.size(), camera_idxs.size());
  THROW_CHECK_EQ(cams_from_rig.size(), cameras->size());
  THROW_CHECK_GE(*std::min_element(camera_idxs.begin(), camera_idxs.end()), 0);
  THROW_CHECK_LT(*std::max_element(camera_idxs.begin(), camera_idxs.end()),
                 cameras->size());
  options.Check();

  const auto loss_function =
      std::make_unique<ceres::CauchyLoss>(options.loss_function_scale);

  std::vector<double*> cameras_params_data(cameras->size());
  for (size_t i = 0; i < cameras->size(); i++) {
    cameras_params_data[i] = cameras->at(i).params.data();
  }
  std::vector<size_t> camera_counts(cameras->size(), 0);

  // Cost function assumes unit quaternion.
  rig_from_world->rotation().normalize();

  std::vector<Eigen::Vector3d> point3D_params = points3D;
  std::vector<Eigen::Vector7d> cam_from_rig_params(cams_from_rig.size());
  for (size_t i = 0; i < cams_from_rig.size(); ++i) {
    cam_from_rig_params[i] = cams_from_rig[i].params;
  }

  ceres::Problem::Options problem_options;
  problem_options.loss_function_ownership = ceres::DO_NOT_TAKE_OWNERSHIP;
  ceres::Problem problem(problem_options);

  for (size_t i = 0; i < points2D.size(); ++i) {
    // Skip outlier observations
    if (!inlier_mask[i]) {
      continue;
    }
    const size_t camera_idx = camera_idxs[i];
    camera_counts[camera_idx] += 1;

    problem.AddResidualBlock(
        CreateCameraCostFunction<RigReprojErrorCostFunctor>(
            cameras->at(camera_idx).model_id, points2D[i]),
        loss_function.get(),
        point3D_params[i].data(),
        cam_from_rig_params[camera_idx].data(),
        rig_from_world->params.data(),
        cameras_params_data[camera_idx]);
    problem.SetParameterBlockConstant(point3D_params[i].data());
  }

  if (options.use_position_prior) {
    problem.AddResidualBlock(
        CovarianceWeightedCostFunctor<AbsolutePosePositionPriorCostFunctor>::
            Create(options.position_prior_covariance,
                   options.position_prior_in_world),
        nullptr,
        rig_from_world->params.data());
  }

  if (problem.NumResiduals() > 0) {
    // Camera parameterization.
    for (size_t i = 0; i < cameras->size(); i++) {
      if (camera_counts[i] == 0) continue;
      Camera& camera = cameras->at(i);

      // We don't optimize the rig parameters (it's likely under-constrained).
      problem.SetParameterBlockConstant(cam_from_rig_params[i].data());

      if (!options.refine_focal_length && !options.refine_extra_params) {
        problem.SetParameterBlockConstant(camera.params.data());
      } else {
        // Always set the principal point as fixed.
        std::vector<int> const_camera_params;
        const span<const size_t> principal_point_idxs =
            camera.PrincipalPointIdxs();
        const_camera_params.insert(const_camera_params.end(),
                                   principal_point_idxs.begin(),
                                   principal_point_idxs.end());

        if (!options.refine_focal_length) {
          const span<const size_t> focal_length_idxs = camera.FocalLengthIdxs();
          const_camera_params.insert(const_camera_params.end(),
                                     focal_length_idxs.begin(),
                                     focal_length_idxs.end());
        }

        if (!options.refine_extra_params) {
          const span<const size_t> extra_params_idxs = camera.ExtraParamsIdxs();
          const_camera_params.insert(const_camera_params.end(),
                                     extra_params_idxs.begin(),
                                     extra_params_idxs.end());
        }

        if (const_camera_params.size() == camera.params.size()) {
          problem.SetParameterBlockConstant(camera.params.data());
        } else {
          SetManifold(
              &problem,
              camera.params.data(),
              CreateSubsetManifold(camera.params.size(), const_camera_params));
        }
      }
    }
  }

  ceres::Solver::Options solver_options;
  solver_options.gradient_tolerance = options.gradient_tolerance;
  solver_options.max_num_iterations = options.max_num_iterations;
  solver_options.linear_solver_type = ceres::DENSE_QR;
  solver_options.logging_type = ceres::LoggingType::SILENT;

  // The overhead of creating threads is too large.
  solver_options.num_threads = 1;
#if CERES_VERSION_MAJOR < 2
  solver_options.num_linear_solver_threads = 1;
#endif  // CERES_VERSION_MAJOR

  ceres::Solver::Summary summary;
  ceres::Solve(solver_options, &problem, &summary);

  if (options.print_summary || VLOG_IS_ON(1)) {
    PrintSolverSummary(summary, "Generalized pose refinement report");
  }

  if (problem.NumResiduals() > 0 && rig_from_world_cov != nullptr) {
    ceres::Covariance::Options options;
    ceres::Covariance covariance(options);
    std::vector<const double*> parameter_blocks = {
        rig_from_world->params.data()};
    if (!covariance.Compute(parameter_blocks, &problem)) {
      return false;
    }
    covariance.GetCovarianceMatrixInTangentSpace(parameter_blocks,
                                                 rig_from_world_cov->data());
  }

  return summary.IsSolutionUsable();
}

bool EstimateStructureLessAbsolutePose(
    const StructureLessAbsolutePoseEstimationOptions& options,
    const std::vector<Eigen::Vector2d>& query_points2D,
    const std::vector<Eigen::Vector2d>& world_points2D,
    const std::vector<size_t>& world_camera_idxs,
    const std::vector<Rigid3d>& world_cams_from_world,
    const std::vector<Camera>& world_cameras,
    const Camera& query_camera,
    Rigid3d* query_cam_from_world,
    size_t* num_inliers,
    std::vector<char>* inlier_mask) {
  THROW_CHECK_EQ(world_points2D.size(), query_points2D.size());
  THROW_CHECK_EQ(world_points2D.size(), world_camera_idxs.size());
  THROW_CHECK_EQ(world_cams_from_world.size(), world_cameras.size());
  ThrowCheckCameras(world_camera_idxs, world_cams_from_world, world_cameras);
  options.Check();

  if (IsPanoramicRig(world_camera_idxs, world_cams_from_world)) {
    return false;
  }

  const size_t num_points = world_points2D.size();
  std::vector<GRNPObservation> world_obs(num_points);
  std::vector<GRNPObservation> query_obs(num_points);
  for (size_t i = 0; i < num_points; ++i) {
    const size_t world_camera_idx = world_camera_idxs[i];
    world_obs[i].cam_from_rig = world_cams_from_world[world_camera_idx];
    if (const std::optional<Eigen::Vector2d> world_cam_point =
            world_cameras[world_camera_idx].CamFromImg(world_points2D[i]);
        world_cam_point.has_value()) {
      world_obs[i].ray_in_cam = world_cam_point->homogeneous().normalized();
    } else {
      world_obs[i].ray_in_cam.setZero();
    }

    query_obs[i].cam_from_rig = Rigid3d();
    if (const std::optional<Eigen::Vector2d> query_cam_point =
            query_camera.CamFromImg(query_points2D[i]);
        query_cam_point.has_value()) {
      query_obs[i].ray_in_cam = query_cam_point->homogeneous().normalized();
    } else {
      query_obs[i].ray_in_cam.setZero();
    }
  }

  auto custom_ransac_options = options.ransac_options;
  custom_ransac_options.max_error =
      query_camera.CamFromImgThreshold(options.ransac_options.max_error);
  LORANSAC<GR6PEstimator, GR8PEstimator> ransac(custom_ransac_options);
  auto report = ransac.Estimate(world_obs, query_obs);
  if (!report.success) {
    return false;
  }

  *query_cam_from_world = report.model;
  *num_inliers = report.support.num_inliers;
  *inlier_mask = std::move(report.inlier_mask);

  return true;
}

}  // namespace colmap
