text: // Copyright (c) 2018, 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.
  //
  // Author: Johannes L. Schoenberger (jsch-at-demuc-dot-de)
  
  #ifndef COLMAP_SRC_RETRIEVAL_INVERTED_INDEX_H_
  #define COLMAP_SRC_RETRIEVAL_INVERTED_INDEX_H_
  
  #include <algorithm>
  #include <bitset>
  #include <cstdint>
  #include <fstream>
  #include <unordered_map>
  #include <unordered_set>
  #include <vector>
  
  #include <Eigen/Core>
  #include <Eigen/Dense>
  
  #include "retrieval/inverted_file.h"
  #include "util/alignment.h"
  #include "util/random.h"
  
  namespace colmap {
  namespace retrieval {
  
  // Implements an inverted index system. The template parameter is the length of
  // the binary vectors in the Hamming Embedding.
  // This class is based on an original implementation by Torsten Sattler.
  template <typename kDescType, int kDescDim, int kEmbeddingDim>
  class InvertedIndex {
   public:
    const static int kInvalidWordId;
    typedef Eigen::Matrix<kDescType, Eigen::Dynamic, kDescDim, Eigen::RowMajor>
        DescType;
    typedef typename InvertedFile<kEmbeddingDim>::EntryType EntryType;
    typedef typename InvertedFile<kEmbeddingDim>::GeomType GeomType;
    typedef Eigen::Matrix<float, Eigen::Dynamic, kDescDim> ProjMatrixType;
    typedef Eigen::VectorXf ProjDescType;
  
    InvertedIndex();
  
    // The number of visual words in the index.
    int NumVisualWords() const;
  
    // Initializes the inverted index with num_words empty inverted files.
    void Initialize(const int num_words);
  
    // Finalizes the inverted index by sorting each inverted file such that all
    // entries are in ascending order of image ids.
    void Finalize();
  
    // Generate projection matrix for Hamming embedding.
    void GenerateHammingEmbeddingProjection();
  
    // Compute Hamming embedding thresholds from a set of descriptors with
    // given visual word identifies.
    void ComputeHammingEmbedding(const DescType& descriptors,
                                 const Eigen::VectorXi& word_ids);
  
    // Add single entry to the index.
    void AddEntry(const int image_id, const int word_id,
                  typename DescType::Index feature_idx,
                  const DescType& descriptor, const GeomType& geometry);
  
    // Clear all index entries.
    void ClearEntries();
  
    // Query the inverted file and return a list of sorted images.
    void Query(const DescType& descriptors, const Eigen::MatrixXi& word_ids,
               std::vector<ImageScore>* image_scores) const;
  
    void ConvertToBinaryDescriptor(
        const int word_id, const DescType& descriptor,
        std::bitset<kEmbeddingDim>* binary_descriptor) const;
  
    float GetIDFWeight(const int word_id) const;
  
    void FindMatches(const int word_id, const std::unordered_set<int>& image_ids,
                     std::vector<const EntryType*>* matches) const;
  
    // Compute the self-similarity for the image.
    float ComputeSelfSimilarity(const Eigen::MatrixXi& word_ids) const;
  
    // Get the identifiers of all indexed images.
    void GetImageIds(std::unordered_set<int>* image_ids) const;
  
    // Read/write the inverted index from/to a binary file.
    void Read(std::ifstream* ifs);
    void Write(std::ofstream* ofs) const;
  
   private:
    void ComputeWeightsAndNormalizationConstants();
  
    // The individual inverted indices.
    std::vector<InvertedFile<kEmbeddingDim>,
                Eigen::aligned_allocator<InvertedFile<kEmbeddingDim>>>
        inverted_files_;
  
    // For each image in the database, a normalization factor to be used to
    // normalize the votes.
    std::unordered_map<int, float> normalization_constants_;
  
    // The projection matrix used to project SIFT descriptors.
    ProjMatrixType proj_matrix_;
  };
  
  ////////////////////////////////////////////////////////////////////////////////
  // Implementation
  ////////////////////////////////////////////////////////////////////////////////
  
  template <typename kDescType, int kDescDim, int kEmbeddingDim>
  const int InvertedIndex<kDescType, kDescDim, kEmbeddingDim>::kInvalidWordId =
      std::numeric_limits<int>::max();
  
  template <typename kDescType, int kDescDim, int kEmbeddingDim>
  InvertedIndex<kDescType, kDescDim, kEmbeddingDim>::InvertedIndex() {
    proj_matrix_.resize(kEmbeddingDim, kDescDim);
    proj_matrix_.setIdentity();
  }
  
  template <typename kDescType, int kDescDim, int kEmbeddingDim>
  int InvertedIndex<kDescType, kDescDim, kEmbeddingDim>::NumVisualWords() const {
    return static_cast<int>(inverted_files_.size());
  }
  
  template <typename kDescType, int kDescDim, int kEmbeddingDim>
  void InvertedIndex<kDescType, kDescDim, kEmbeddingDim>::Initialize(
      const int num_words) {
    CHECK_GT(num_words, 0);
    inverted_files_.resize(num_words);
    for (auto& inverted_file : inverted_files_) {
      inverted_file.Reset();
    }
  }
  
  template <typename kDescType, int kDescDim, int kEmbeddingDim>
  void InvertedIndex<kDescType, kDescDim, kEmbeddingDim>::Finalize() {
    CHECK_GT(NumVisualWords(), 0);
  
    for (auto& inverted_file : inverted_files_) {
      inverted_file.SortEntries();
    }
  
    ComputeWeightsAndNormalizationConstants();
  }
  
  template <typename kDescType, int kDescDim, int kEmbeddingDim>
  void InvertedIndex<kDescType, kDescDim,
                     kEmbeddingDim>::GenerateHammingEmbeddingProjection() {
    Eigen::MatrixXf random_matrix(kDescDim, kDescDim);
    for (Eigen::MatrixXf::Index i = 0; i < random_matrix.size(); ++i) {
      random_matrix(i) = RandomGaussian(0.0f, 1.0f);
    }
    const Eigen::MatrixXf Q = random_matrix.colPivHouseholderQr().matrixQ();
    proj_matrix_ = Q.topRows<kEmbeddingDim>();
  }
  
  template <typename kDescType, int kDescDim, int kEmbeddingDim>
  void InvertedIndex<kDescType, kDescDim, kEmbeddingDim>::ComputeHammingEmbedding(
      const DescType& descriptors, const Eigen::VectorXi& word_ids) {
    CHECK_EQ(descriptors.rows(), word_ids.rows());
    CHECK_EQ(descriptors.cols(), kDescDim);
  
    // Skip every inverted file with less than kMinEntries entries.
    const size_t kMinEntries = 5;
  
    // Determines for each word the corresponding descriptors.
    std::vector<std::vector<int>> indices_per_word(NumVisualWords());
    for (Eigen::MatrixXi::Index i = 0; i < word_ids.rows(); ++i) {
      indices_per_word.at(word_ids(i)).push_back(i);
    }
  
    // For each word, learn the Hamming embedding threshold and the local
    // descriptor space densities.
    for (int i = 0; i < NumVisualWords(); ++i) {
      const auto& indices = indices_per_word[i];
      if (indices.size() < kMinEntries) {
        continue;
      }
  
      Eigen::Matrix<float, Eigen::Dynamic, kEmbeddingDim> proj_desc(
          indices.size(), kEmbeddingDim);
      for (size_t j = 0; j < indices.size(); ++j) {
        proj_desc.row(j) =
            proj_matrix_ *
            descriptors.row(indices[j]).transpose().template cast<float>();
      }
  
      inverted_files_[i].ComputeHammingEmbedding(proj_desc);
    }
  }
  
  template <typename kDescType, int kDescDim, int kEmbeddingDim>
  void InvertedIndex<kDescType, kDescDim, kEmbeddingDim>::AddEntry(
      const int image_id, const int word_id, typename DescType::Index feature_idx,
      const DescType& descriptor, const GeomType& geometry) {
    CHECK_EQ(descriptor.size(), kDescDim);
    const ProjDescType proj_desc =
        proj_matrix_ * descriptor.transpose().template cast<float>();
    inverted_files_.at(word_id).AddEntry(image_id, feature_idx, proj_desc,
                                         geometry);
  }
  
  template <typename kDescType, int kDescDim, int kEmbeddingDim>
  void InvertedIndex<kDescType, kDescDim, kEmbeddingDim>::ClearEntries() {
    for (auto& inverted_file : inverted_files_) {
      inverted_file.ClearEntries();
    }
  }
  
  template <typename kDescType, int kDescDim, int kEmbeddingDim>
  void InvertedIndex<kDescType, kDescDim, kEmbeddingDim>::Query(
      const DescType& descriptors, const Eigen::MatrixXi& word_ids,
      std::vector<ImageScore>* image_scores) const {
    CHECK_EQ(descriptors.cols(), kDescDim);
  
    image_scores->clear();
  
    // Computes the self-similarity score for the query image.
    const float self_similarity = ComputeSelfSimilarity(word_ids);
    float normalization_weight = 1.0f;
    if (self_similarity > 0.0f) {
      normalization_weight = 1.0f / std::sqrt(self_similarity);
    }
  
    std::unordered_map<int, int> score_map;
    std::vector<ImageScore> inverted_file_scores;
  
    for (typename DescType::Index i = 0; i < descriptors.rows(); ++i) {
      const ProjDescType proj_descriptor =
          proj_matrix_ * descriptors.row(i).transpose().template cast<float>();
      for (Eigen::MatrixXi::Index n = 0; n < word_ids.cols(); ++n) {
        const int word_id = word_ids(i, n);
        if (word_id == kInvalidWordId) {
          continue;
        }
  
        inverted_files_.at(word_id).ScoreFeature(proj_descriptor,
                                                 &inverted_file_scores);
  
        for (const ImageScore& score : inverted_file_scores) {
          const auto score_map_it = score_map.find(score.image_id);
          if (score_map_it == score_map.end()) {
            // Image not found in another inverted file.
            score_map.emplace(score.image_id,
                              static_cast<int>(image_scores->size()));
            image_scores->push_back(score);
          } else {
            // Image already found in another inverted file, so accumulate.
            (*image_scores).at(score_map_it->second).score += score.score;
          }
        }
      }
    }
  
    // Normalization.
    for (ImageScore& score : *image_scores) {
      score.score *=
          normalization_weight * normalization_constants_.at(score.image_id);
    }
  }
  
  template <typename kDescType, int kDescDim, int kEmbeddingDim>
  void InvertedIndex<kDescType, kDescDim, kEmbeddingDim>::
      ConvertToBinaryDescriptor(
          const int word_id, const DescType& descriptor,
          std::bitset<kEmbeddingDim>* binary_descriptor) const {
    const ProjDescType proj_desc =
        proj_matrix_ * descriptor.transpose().template cast<float>();
    inverted_files_.at(word_id).ConvertToBinaryDescriptor(proj_desc,
                                                          binary_descriptor);
  }
  
  template <typename kDescType, int kDescDim, int kEmbeddingDim>
  float InvertedIndex<kDescType, kDescDim, kEmbeddingDim>::GetIDFWeight(
      const int word_id) const {
    return inverted_files_.at(word_id).IDFWeight();
  }
  
  template <typename kDescType, int kDescDim, int kEmbeddingDim>
  void InvertedIndex<kDescType, kDescDim, kEmbeddingDim>::FindMatches(
      const int word_id, const std::unordered_set<int>& image_ids,
      std::vector<const EntryType*>* matches) const {
    matches->clear();
    const auto& entries = inverted_files_.at(word_id).GetEntries();
    for (const auto& entry : entries) {
      if (image_ids.count(entry.image_id)) {
        matches->emplace_back(&entry);
      }
    }
  }
  
  template <typename kDescType, int kDescDim, int kEmbeddingDim>
  float InvertedIndex<kDescType, kDescDim, kEmbeddingDim>::ComputeSelfSimilarity(
      const Eigen::MatrixXi& word_ids) const {
    double self_similarity = 0.0;
    for (Eigen::MatrixXi::Index i = 0; i < word_ids.size(); ++i) {
      const int word_id = word_ids(i);
      if (word_id != kInvalidWordId) {
        const auto& inverted_file = inverted_files_.at(word_id);
        self_similarity += inverted_file.IDFWeight() * inverted_file.IDFWeight();
      }
    }
    return static_cast<float>(self_similarity);
  }
  
  template <typename kDescType, int kDescDim, int kEmbeddingDim>
  void InvertedIndex<kDescType, kDescDim, kEmbeddingDim>::GetImageIds(
      std::unordered_set<int>* image_ids) const {
    for (const auto& inverted_file : inverted_files_) {
      inverted_file.GetImageIds(image_ids);
    }
  }
  
  template <typename kDescType, int kDescDim, int kEmbeddingDim>
  void InvertedIndex<kDescType, kDescDim, kEmbeddingDim>::Read(
      std::ifstream* ifs) {
    CHECK(ifs->is_open());
  
    int32_t num_words = 0;
    ifs->read(reinterpret_cast<char*>(&num_words), sizeof(int32_t));
    CHECK_GT(num_words, 0);
  
    Initialize(num_words);
  
    int32_t N_t = 0;
    ifs->read(reinterpret_cast<char*>(&N_t), sizeof(int32_t));
    CHECK_EQ(N_t, kEmbeddingDim)
        << "The length of the binary strings should be " << kEmbeddingDim
        << " but is " << N_t << ". The indices are not compatible!";
  
    for (int i = 0; i < kEmbeddingDim; ++i) {
      for (int j = 0; j < kDescDim; ++j) {
        ifs->read(reinterpret_cast<char*>(&proj_matrix_(i, j)), sizeof(float));
      }
    }
  
    for (auto& inverted_file : inverted_files_) {
      inverted_file.Read(ifs);
    }
  
    int32_t num_images = 0;
    ifs->read(reinterpret_cast<char*>(&num_images), sizeof(int32_t));
    CHECK_GE(num_images, 0);
  
    normalization_constants_.clear();
    normalization_constants_.reserve(num_images);
    for (int32_t i = 0; i < num_images; ++i) {
      int image_id;
      float value;
      ifs->read(reinterpret_cast<char*>(&image_id), sizeof(int));
      ifs->read(reinterpret_cast<char*>(&value), sizeof(float));
      normalization_constants_[image_id] = value;
    }
  }
  
  template <typename kDescType, int kDescDim, int kEmbeddingDim>
  void InvertedIndex<kDescType, kDescDim, kEmbeddingDim>::Write(
      std::ofstream* ofs) const {
    CHECK(ofs->is_open());
  
    int32_t num_words = static_cast<int32_t>(NumVisualWords());
    ofs->write(reinterpret_cast<const char*>(&num_words), sizeof(int32_t));
    CHECK_GT(num_words, 0);
  
    const int32_t N_t = static_cast<int32_t>(kEmbeddingDim);
    ofs->write(reinterpret_cast<const char*>(&N_t), sizeof(int32_t));
  
    for (int i = 0; i < kEmbeddingDim; ++i) {
      for (int j = 0; j < kDescDim; ++j) {
        ofs->write(reinterpret_cast<const char*>(&proj_matrix_(i, j)),
                   sizeof(float));
      }
    }
  
    for (const auto& inverted_file : inverted_files_) {
      inverted_file.Write(ofs);
    }
  
    const int32_t num_images = normalization_constants_.size();
    ofs->write(reinterpret_cast<const char*>(&num_images), sizeof(int32_t));
  
    for (const auto& constant : normalization_constants_) {
      ofs->write(reinterpret_cast<const char*>(&constant.first), sizeof(int));
      ofs->write(reinterpret_cast<const char*>(&constant.second), sizeof(float));
    }
  }
  
  template <typename kDescType, int kDescDim, int kEmbeddingDim>
  void InvertedIndex<kDescType, kDescDim,
                     kEmbeddingDim>::ComputeWeightsAndNormalizationConstants() {
    std::unordered_set<int> image_ids;
    GetImageIds(&image_ids);
  
    for (auto& inverted_file : inverted_files_) {
      inverted_file.ComputeIDFWeight(image_ids.size());
    }
  
    std::unordered_map<int, double> self_similarities(image_ids.size());
    for (const auto& inverted_file : inverted_files_) {
      inverted_file.ComputeImageSelfSimilarities(&self_similarities);
    }
  
    normalization_constants_.clear();
    normalization_constants_.reserve(image_ids.size());
    for (const auto& self_similarity : self_similarities) {
      if (self_similarity.second > 0.0) {
        normalization_constants_[self_similarity.first] =
            static_cast<float>(1.0 / std::sqrt(self_similarity.second));
      } else {
        normalization_constants_[self_similarity.first] = 0.0f;
      }
    }
  }
  
  }  // namespace retrieval
  }  // namespace colmap
  
  #endif  // COLMAP_SRC_RETRIEVAL_INVERTED_INDEX_H_
  