// Copyright 2015 Google Inc. All Rights Reserved.
//
// 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.
//

// Author: ericv@google.com (Eric Veach)

#ifndef S2_S2POINT_INDEX_H_
#define S2_S2POINT_INDEX_H_

#include <cstddef>
#include <tuple>
#include <type_traits>
#include <utility>

#include "absl/container/btree_map.h"
#include "absl/log/absl_check.h"

#include "s2/_fp_contract_off.h"  // IWYU pragma: keep
#include "s2/s2cell_id.h"
#include "s2/s2cell_iterator.h"
#include "s2/s2point.h"

namespace s2internal {
// Hack to expose bytes_used.
template <typename Key, typename Value>
class BTreeMultimap : public absl::btree_multimap<Key, Value> {
 public:
  size_t bytes_used() const { return this->tree_.bytes_used(); }
};
}  // namespace s2internal


// S2PointIndex maintains an index of points sorted by leaf S2CellId.  Each
// point can optionally store auxiliary data such as an integer or pointer.
// This can be used to map results back to client data structures.
//
// The class supports adding or removing points dynamically, and provides a
// seekable iterator interface for navigating the index.
//
// You can use this class in conjunction with S2ClosestPointQuery to find the
// closest index points to a given query point.  For example:
//
// void Test(const vector<S2Point>& index_points,
//           const vector<S2Point>& target_points) {
//   // The template argument allows auxiliary data to be attached to each
//   // point (in this case, the array index).
//   S2PointIndex<int> index;
//   for (int i = 0; i < index_points.size(); ++i) {
//     index.Add(index_points[i], i);
//   }
//   S2ClosestPointQuery<int> query(&index);
//   query.mutable_options()->set_max_results(5);
//   for (const S2Point& target_point : target_points) {
//     S2ClosestPointQueryPointTarget target(target_point);
//     for (const auto& result : query.FindClosestPoints(&target)) {
//       // The Result class contains the following methods:
//       //   distance() is the distance to the target.
//       //   point() is the indexed point.
//       //   data() is the auxiliary data.
//       DoSomething(target_point, result);
//     }
//   }
// }
//
// The Data argument defaults to an empty class, which uses no additional
// space beyond the S2Point itself.  In this case the Data argument is
// required.  For example:
//
//   S2PointIndex<> index;
//   index.Add(point);
//
// Points can be added or removed from the index at any time by calling Add()
// or Remove().  However when the index is modified, you must call Init() on
// each iterator before using it again (or simply create a new iterator).
//
//   index.Add(new_point, 123456);
//   it.Init(&index);
//   it.Seek(target.range_min());
//
// You can also access the index directly using the iterator interface.  For
// example, here is how to iterate through all the points in a given S2CellId
// "target_id":
//
//   S2PointIndex<int>::Iterator it(&index);
//   it.Seek(target_id.range_min());
//   for (; !it.done() && it.id() <= target_id.range_max(); it.Next()) {
//     DoSomething(it.id(), it.point(), it.data());
//   }
//
// TODO(ericv): Consider adding an S2PointIndexRegion class, which could be
// used to efficiently compute coverings of a collection of S2Points.
//
// REQUIRES: "Data" has default and copy constructors.
// REQUIRES: "Data" has operator== and operator<.
template <class Data = std::tuple<> /*empty class*/>
class S2PointIndex {
 public:
  // PointData is essentially std::pair with named fields.  It stores an
  // S2Point and its associated data, taking advantage of the "empty base
  // optimization" to ensure that no extra space is used when Data is empty.
  class PointData {
   public:
    PointData() = default;  // Needed by STL
    PointData(const S2Point& point, const Data& data) : tuple_(point, data) {}

    const S2Point& point() const { return std::get<0>(tuple_); }
    const Data& data() const { return std::get<1>(tuple_); }

    friend bool operator==(const PointData& x, const PointData& y) {
      return x.tuple_ == y.tuple_;
    }
    friend bool operator<(const PointData& x, const PointData& y) {
      return x.tuple_ < y.tuple_;
    }

   private:
    // Note that std::tuple has special template magic to ensure that Data
    // doesn't take up any space when it is empty.  (This is not true if you
    // simply declare a member of type Data.)
    std::tuple<S2Point, Data> tuple_;
  };

  // Default constructor.
  // TODO(b/252809194): Move definition back to .cc file.
  S2PointIndex() = default;

  // Returns the number of points in the index.
  int num_points() const;

  // Adds the given point to the index.  Invalidates all iterators.
  void Add(const S2Point& point, const Data& data);
  void Add(const PointData& point_data);

  // Convenience function for the case when Data is an empty class.
  void Add(const S2Point& point);

  // Removes the given point from the index.  Both the "point" and "data"
  // fields must match the point to be removed.  Returns false if the given
  // point was not present.  Invalidates all iterators.
  bool Remove(const S2Point& point, const Data& data);
  bool Remove(const PointData& point_data);

  // Convenience function for the case when Data is an empty class.
  void Remove(const S2Point& point);

  // Resets the index to its original empty state.  Invalidates all iterators.
  void Clear();

  // Returns the number of bytes currently occupied by the index.
  size_t SpaceUsed() const;

 private:
  // Defined here because the Iterator class below uses it.
  using Map = s2internal::BTreeMultimap<S2CellId, PointData>;

 public:
  class Iterator final : public S2CellIterator {
   public:
    // Default constructor; must be followed by a call to Init().
    Iterator();

    // Convenience constructor that calls Init().
    explicit Iterator(const S2PointIndex* index);

    // Initializes an iterator for the given S2PointIndex.  If the index is
    // non-empty, the iterator is positioned at the first cell.
    //
    // This method may be called multiple times, e.g. to make an iterator
    // valid again after the index is modified.
    void Init(const S2PointIndex* index);

    // The S2CellId for the current index entry.
    // REQUIRES: !done()
    S2CellId id() const override;

    // The point associated with the current index entry.
    // REQUIRES: !done()
    const S2Point& point() const;

    // The client-supplied data associated with the current index entry.
    // REQUIRES: !done()
    const Data& data() const;

    // The (S2Point, data) pair associated with the current index entry.
    const PointData& point_data() const;

    // Returns true if the iterator is positioned past the last index entry.
    bool done() const override;

    // Positions the iterator at the first index entry (if any).
    void Begin() override;

    // Positions the iterator so that done() is true.
    void Finish() override;

    // Advances the iterator to the next index entry.
    // REQUIRES: !done()
    void Next() override;

    // If the iterator is already positioned at the beginning, returns false.
    // Otherwise positions the iterator at the previous entry and returns true.
    bool Prev() override;

    // Positions the iterator at the first entry with id() >= target, or at the
    // end of the index if no such entry exists.
    void Seek(S2CellId target) override;

    // Positions the iterator at the cell containing target and returns true. If
    // no such cell exists, return false and leave the iterator in an undefined
    // (but valid) state.
    bool Locate(const S2Point& target) override {
      return LocateImpl(*this, target);
    }

    // Let T be the target S2CellId.  If T is contained by some index cell I
    // (including equality), this method positions the iterator at I and returns
    // INDEXED.  Otherwise if T contains one or more (smaller) index cells, it
    // positions the iterator at the first such cell I and returns SUBDIVIDED.
    // Otherwise it returns DISJOINT and leaves the iterator in an undefined
    // (but valid) state.
    S2CellRelation Locate(S2CellId target) override {
      return LocateImpl(*this, target);
    }

   private:
    const Map* map_ = nullptr;
    typename Map::const_iterator iter_, end_;
  };

 private:
  friend class Iterator;
  Map map_;

  S2PointIndex(const S2PointIndex&) = delete;
  void operator=(const S2PointIndex&) = delete;
};


//////////////////   Implementation details follow   ////////////////////

template <class Data>
inline int S2PointIndex<Data>::num_points() const {
  return map_.size();
}

template <class Data>
void S2PointIndex<Data>::Add(const PointData& point_data) {
  S2CellId id(point_data.point());
  map_.insert(std::make_pair(id, point_data));
}

template <class Data>
void S2PointIndex<Data>::Add(const S2Point& point, const Data& data) {
  Add(PointData(point, data));
}

template <class Data>
void S2PointIndex<Data>::Add(const S2Point& point) {
  static_assert(std::is_empty<Data>::value, "Data must be empty");
  Add(point, {});
}

template <class Data>
bool S2PointIndex<Data>::Remove(const PointData& point_data) {
  S2CellId id(point_data.point());
  for (typename Map::iterator it = map_.lower_bound(id), end = map_.end();
       it != end && it->first == id; ++it) {
    if (it->second == point_data) {
      map_.erase(it);
      return true;
    }
  }
  return false;
}

template <class Data>
bool S2PointIndex<Data>::Remove(const S2Point& point, const Data& data) {
  return Remove(PointData(point, data));
}

template <class Data>
void S2PointIndex<Data>::Remove(const S2Point& point) {
  static_assert(std::is_empty<Data>::value, "Data must be empty");
  Remove(point, {});
}

template <class Data>
void S2PointIndex<Data>::Clear() {
  map_.clear();
}

template <class Data>
size_t S2PointIndex<Data>::SpaceUsed() const {
  return sizeof(*this) - sizeof(map_) + map_.bytes_used();
}

template <class Data>
inline S2PointIndex<Data>::Iterator::Iterator() = default;

template <class Data>
inline S2PointIndex<Data>::Iterator::Iterator(
    const S2PointIndex<Data>* index) {
  Init(index);
}

template <class Data>
inline void S2PointIndex<Data>::Iterator::Init(
    const S2PointIndex<Data>* index) {
  map_ = &index->map_;
  iter_ = map_->begin();
  end_ = map_->end();
}

template <class Data>
inline S2CellId S2PointIndex<Data>::Iterator::id() const {
  ABSL_DCHECK(!done());
  return iter_->first;
}

template <class Data>
inline const S2Point& S2PointIndex<Data>::Iterator::point() const {
  ABSL_DCHECK(!done());
  return iter_->second.point();
}

template <class Data>
inline const Data& S2PointIndex<Data>::Iterator::data() const {
  ABSL_DCHECK(!done());
  return iter_->second.data();
}

template <class Data>
inline const typename S2PointIndex<Data>::PointData&
S2PointIndex<Data>::Iterator::point_data() const {
  ABSL_DCHECK(!done());
  return iter_->second;
}

template <class Data>
inline bool S2PointIndex<Data>::Iterator::done() const {
  return iter_ == end_;
}

template <class Data>
inline void S2PointIndex<Data>::Iterator::Begin() {
  iter_ = map_->begin();
}

template <class Data>
inline void S2PointIndex<Data>::Iterator::Finish() {
  iter_ = end_;
}

template <class Data>
inline void S2PointIndex<Data>::Iterator::Next() {
  ABSL_DCHECK(!done());
  ++iter_;
}

template <class Data>
inline bool S2PointIndex<Data>::Iterator::Prev() {
  if (iter_ == map_->begin()) return false;
  --iter_;
  return true;
}

template <class Data>
inline void S2PointIndex<Data>::Iterator::Seek(S2CellId target) {
  iter_ = map_->lower_bound(target);
}

#endif  // S2_S2POINT_INDEX_H_
