/*
 * Copyright 2016 The Cartographer Authors
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *      http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

// Implements the loop closure method called Sparse Pose Adjustment (SPA) from
// Konolige, Kurt, et al. "Efficient sparse pose adjustment for 2d mapping."
// Intelligent Robots and Systems (IROS), 2010 IEEE/RSJ International Conference
// on (pp. 22--29). IEEE, 2010.
//
// It is extended for submapping in 3D:
// Each scan has been matched against one or more submaps (adding a constraint
// for each match), both poses of scans and of submaps are to be optimized.
// All constraints are between a submap i and a scan j.

#ifndef CARTOGRAPHER_MAPPING_3D_SPARSE_POSE_GRAPH_H_
#define CARTOGRAPHER_MAPPING_3D_SPARSE_POSE_GRAPH_H_

#include <deque>
#include <functional>
#include <limits>
#include <map>
#include <unordered_map>
#include <vector>

#include "Eigen/Core"
#include "Eigen/Geometry"
#include "cartographer/common/fixed_ratio_sampler.h"
#include "cartographer/common/mutex.h"
#include "cartographer/common/thread_pool.h"
#include "cartographer/common/time.h"
#include "cartographer/kalman_filter/pose_tracker.h"
#include "cartographer/mapping/sparse_pose_graph.h"
#include "cartographer/mapping/trajectory_connectivity.h"
#include "cartographer/mapping_3d/sparse_pose_graph/constraint_builder.h"
#include "cartographer/mapping_3d/sparse_pose_graph/optimization_problem.h"
#include "cartographer/mapping_3d/submaps.h"
#include "cartographer/sensor/point_cloud.h"
#include "cartographer/transform/rigid_transform.h"
#include "cartographer/transform/transform.h"

namespace cartographer {
namespace mapping_3d {

// Implements the SPA loop closure method.
class SparsePoseGraph : public mapping::SparsePoseGraph {
 public:
  SparsePoseGraph(const mapping::proto::SparsePoseGraphOptions& options,
                  common::ThreadPool* thread_pool);
  ~SparsePoseGraph() override;

  SparsePoseGraph(const SparsePoseGraph&) = delete;
  SparsePoseGraph& operator=(const SparsePoseGraph&) = delete;

  // Adds a new 'range_data_in_tracking' observation at 'time', and a 'pose'
  // that will later be optimized. The 'pose' was determined by scan matching
  // against the 'matching_submap' and the scan was inserted into the
  // 'insertion_submaps'.
  void AddScan(common::Time time,
               const sensor::RangeData& range_data_in_tracking,
               const transform::Rigid3d& pose,
               const kalman_filter::PoseCovariance& pose_covariance,
               const Submaps* trajectory, const Submap* matching_submap,
               const std::vector<const Submap*>& insertion_submaps)
      EXCLUDES(mutex_);

  // The index into the vector of trajectory nodes as used with
  // GetTrajectoryNodes() for the next node added with AddScan() is returned.
  int GetNextTrajectoryNodeIndex() EXCLUDES(mutex_);

  // Adds new IMU data to be used in the optimization.
  void AddImuData(const mapping::Submaps* trajectory, common::Time time,
                  const Eigen::Vector3d& linear_acceleration,
                  const Eigen::Vector3d& angular_velocity);

  void RunFinalOptimization() override;
  std::vector<std::vector<const mapping::Submaps*>> GetConnectedTrajectories()
      override;
  std::vector<transform::Rigid3d> GetSubmapTransforms(
      const mapping::Submaps* trajectory) EXCLUDES(mutex_) override;
  transform::Rigid3d GetLocalToGlobalTransform(const mapping::Submaps* submaps)
      EXCLUDES(mutex_) override;
  std::vector<mapping::TrajectoryNode> GetTrajectoryNodes() override
      EXCLUDES(mutex_);
  std::vector<SubmapState> GetSubmapStates() override;
  std::vector<Constraint> constraints() override;

 private:
  // This is 'mapping::SubmapState', but with the 3D versions of 'submap' and
  // 'trajectory'.
  struct SubmapState {
    const Submap* submap = nullptr;

    // Indices of the scans that were inserted into this map together with
    // constraints for them. They are not to be matched again when this submap
    // becomes 'finished'.
    std::set<int> scan_indices;

    // Whether in the current state of the background thread this submap is
    // finished. When this transitions to true, all scans are tried to match
    // against this submap. Likewise, all new scans are matched against submaps
    // which are finished.
    bool finished = false;

    // The trajectory to which this SubmapState belongs.
    const Submaps* trajectory = nullptr;
  };

  // Handles a new work item.
  void AddWorkItem(std::function<void()> work_item) REQUIRES(mutex_);

  int GetSubmapIndex(const mapping::Submap* submap) const REQUIRES(mutex_) {
    const auto iterator = submap_indices_.find(submap);
    CHECK(iterator != submap_indices_.end());
    return iterator->second;
  }

  // Grows 'submap_transforms_' to have an entry for every element of 'submaps'.
  void GrowSubmapTransformsAsNeeded(const std::vector<const Submap*>& submaps)
      REQUIRES(mutex_);

  // Adds constraints for a scan, and starts scan matching in the background.
  void ComputeConstraintsForScan(
      int scan_index, const Submap* matching_submap,
      std::vector<const Submap*> insertion_submaps,
      const Submap* finished_submap, const transform::Rigid3d& pose,
      const kalman_filter::PoseCovariance& covariance) REQUIRES(mutex_);

  // Computes constraints for a scan and submap pair.
  void ComputeConstraint(const int scan_index, const int submap_index)
      REQUIRES(mutex_);

  // Adds constraints for older scans whenever a new submap is finished.
  void ComputeConstraintsForOldScans(const Submap* submap) REQUIRES(mutex_);

  // Registers the callback to run the optimization once all constraints have
  // been computed, that will also do all work that queue up in 'scan_queue_'.
  void HandleScanQueue() REQUIRES(mutex_);

  // Waits until we caught up (i.e. nothing is waiting to be scheduled), and
  // all computations have finished.
  void WaitForAllComputations() EXCLUDES(mutex_);

  // Runs the optimization. Callers have to make sure, that there is only one
  // optimization being run at a time.
  void RunOptimization() EXCLUDES(mutex_);

  // Adds extrapolated transforms, so that there are transforms for all submaps.
  std::vector<transform::Rigid3d> ExtrapolateSubmapTransforms(
      const std::vector<transform::Rigid3d>& submap_transforms,
      const mapping::Submaps* trajectory) const REQUIRES(mutex_);

  const mapping::proto::SparsePoseGraphOptions options_;
  common::Mutex mutex_;

  // If it exists, further scans must be added to this queue, and will be
  // considered later.
  std::unique_ptr<std::deque<std::function<void()>>> scan_queue_
      GUARDED_BY(mutex_);

  // How our various trajectories are related.
  mapping::TrajectoryConnectivity trajectory_connectivity_ GUARDED_BY(mutex_);

  // We globally localize a fraction of the scans from each trajectory.
  std::unordered_map<const mapping::Submaps*,
                     std::unique_ptr<common::FixedRatioSampler>>
      global_localization_samplers_ GUARDED_BY(mutex_);

  // Number of scans added since last loop closure.
  int num_scans_since_last_loop_closure_ GUARDED_BY(mutex_) = 0;

  // Whether the optimization has to be run before more data is added.
  bool run_loop_closure_ GUARDED_BY(mutex_) = false;

  // Current optimization problem.
  sparse_pose_graph::OptimizationProblem optimization_problem_;
  sparse_pose_graph::ConstraintBuilder constraint_builder_ GUARDED_BY(mutex_);
  std::vector<Constraint> constraints_;
  std::vector<transform::Rigid3d> submap_transforms_;  // (map <- submap)

  // Submaps get assigned an index and state as soon as they are seen, even
  // before they take part in the background computations.
  std::map<const mapping::Submap*, int> submap_indices_ GUARDED_BY(mutex_);
  std::vector<SubmapState> submap_states_ GUARDED_BY(mutex_);

  // Connectivity structure of our trajectories.
  std::vector<std::vector<const mapping::Submaps*>> connected_components_;
  // Trajectory to connected component ID.
  std::map<const mapping::Submaps*, size_t> reverse_connected_components_;

  // Data that are currently being shown.
  //
  // Deque to keep references valid for the background computation when adding
  // new data.
  std::deque<mapping::TrajectoryNode::ConstantData> constant_node_data_;
  std::vector<mapping::TrajectoryNode> trajectory_nodes_ GUARDED_BY(mutex_);

  // Current submap transforms used for displaying data.
  std::vector<transform::Rigid3d> optimized_submap_transforms_
      GUARDED_BY(mutex_);
};

}  // namespace mapping_3d
}  // namespace cartographer

#endif  // CARTOGRAPHER_MAPPING_3D_SPARSE_POSE_GRAPH_H_