// 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.

#pragma once

#include "colmap/math/math.h"
#include "colmap/retrieval/geometry.h"
#include "colmap/retrieval/inverted_file_entry.h"
#include "colmap/retrieval/utils.h"
#include "colmap/util/eigen_alignment.h"
#include "colmap/util/endian.h"
#include "colmap/util/logging.h"

#include <algorithm>
#include <bitset>
#include <cstdint>
#include <unordered_map>
#include <unordered_set>
#include <vector>

#include <Eigen/Core>

namespace colmap {
namespace retrieval {

// Implements an inverted file, including the ability to compute image scores
// and matches. 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 <int kEmbeddingDim>
class InvertedFile {
 public:
  using DescType = Eigen::VectorXf;
  using GeomType = FeatureGeometry;
  using EntryType = InvertedFileEntry<kEmbeddingDim>;

  enum Status : uint8_t {
    kUnusable = 0x00,
    kHasEmbedding = 0x01,
    kEntriesSorted = 0x02,
    kUsable = 0x03,
  };

  InvertedFile();

  // The number of added entries.
  size_t NumEntries() const;

  // Return all entries in the file.
  const std::vector<EntryType>& Entries() const;

  // Whether the Hamming embedding was computed for this file.
  bool HasHammingEmbedding() const;

  // Whether the entries in this file are sorted.
  bool EntriesSorted() const;

  // Whether this file is usable for scoring, i.e. the entries are sorted and
  // the Hamming embedding has been computed.
  bool IsUsable() const;

  // Adds an inverted file entry given a projected descriptor and its image
  // information stored in an inverted file entry. In particular, this function
  // generates the binary descriptor for the inverted file entry and then stores
  // the entry in the inverted file.
  void AddEntry(int image_id,
                typename DescType::Index feature_idx,
                const DescType& descriptor,
                const GeomType& geometry);

  // Sorts the inverted file entries in ascending order of image ids. This is
  // required for efficient scoring and must be called before ScoreFeature.
  void SortEntries();

  // Clear all entries in this file.
  void ClearEntries();

  // Reset all computed weights/thresholds and clear all entries.
  void Reset();

  // Given a projected descriptor, returns the corresponding binary string.
  void ConvertToBinaryDescriptor(
      const DescType& descriptor,
      std::bitset<kEmbeddingDim>* binary_descriptor) const;

  // Compute the idf-weight for this inverted file.
  void ComputeIDFWeight(int num_total_images);

  // Return the idf-weight of this inverted file.
  float SquaredIDFWeight() const;

  // Given a set of descriptors, learns the thresholds required for the Hamming
  // embedding. Each row in descriptors represents a single descriptor projected
  // into the kEmbeddingDim dimensional space used for Hamming embedding.
  void ComputeHammingEmbedding(
      const Eigen::Matrix<float, Eigen::Dynamic, kEmbeddingDim>& descriptors);

  // Given a query feature, performs inverted file scoring.
  void ScoreFeature(const DescType& descriptor,
                    std::vector<ImageScore>* image_scores) const;

  // Get the identifiers of all indexed images in this file.
  void GetImageIds(std::unordered_set<int>* ids) const;

  // For each image in the inverted file, computes the self-similarity of each
  // image in the file (the part caused by this word) and adds the weight to the
  // entry corresponding to that image. This function is useful to determine the
  // normalization factor for each image that is used during retrieval.
  void ComputeImageSelfSimilarities(
      std::unordered_map<int, double>* self_similarities) const;

  // Read/write the inverted file from/to a binary file.
  void Read(std::istream* in);
  void Write(std::ostream* out) const;

 private:
  // Whether the inverted file is initialized.
  uint8_t status_;

  // The inverse document frequency weight of this inverted file.
  float squared_idf_weight_;

  // The entries of the inverted file system.
  std::vector<EntryType> entries_;

  // The thresholds used for Hamming embedding.
  DescType thresholds_;

  // The functor to derive a voting weight from a Hamming distance.
  static const HammingDistWeightFunctor<kEmbeddingDim>
      hamming_dist_weight_functor_;
};

////////////////////////////////////////////////////////////////////////////////
// Implementation
////////////////////////////////////////////////////////////////////////////////

template <int kEmbeddingDim>
const HammingDistWeightFunctor<kEmbeddingDim>
    InvertedFile<kEmbeddingDim>::hamming_dist_weight_functor_;

template <int kEmbeddingDim>
InvertedFile<kEmbeddingDim>::InvertedFile()
    : status_(kUnusable), squared_idf_weight_(0.0f) {
  static_assert(kEmbeddingDim % 8 == 0,
                "Dimensionality of projected space needs to"
                " be a multiple of 8.");
  static_assert(kEmbeddingDim > 0,
                "Dimensionality of projected space needs to be > 0.");

  thresholds_.resize(kEmbeddingDim);
  thresholds_.setZero();
}

template <int kEmbeddingDim>
size_t InvertedFile<kEmbeddingDim>::NumEntries() const {
  return entries_.size();
}

template <int kEmbeddingDim>
const std::vector<typename InvertedFile<kEmbeddingDim>::EntryType>&
InvertedFile<kEmbeddingDim>::Entries() const {
  return entries_;
}

template <int kEmbeddingDim>
bool InvertedFile<kEmbeddingDim>::HasHammingEmbedding() const {
  return status_ & kHasEmbedding;
}

template <int kEmbeddingDim>
bool InvertedFile<kEmbeddingDim>::EntriesSorted() const {
  return status_ & kEntriesSorted;
}

template <int kEmbeddingDim>
bool InvertedFile<kEmbeddingDim>::IsUsable() const {
  return status_ & kUsable;
}

template <int kEmbeddingDim>
void InvertedFile<kEmbeddingDim>::AddEntry(const int image_id,
                                           typename DescType::Index feature_idx,
                                           const DescType& descriptor,
                                           const GeomType& geometry) {
  THROW_CHECK_GE(image_id, 0);
  THROW_CHECK_EQ(descriptor.size(), kEmbeddingDim);
  EntryType entry;
  entry.image_id = image_id;
  entry.feature_idx = feature_idx;
  entry.geometry = geometry;
  ConvertToBinaryDescriptor(descriptor, &entry.descriptor);
  entries_.push_back(entry);
  status_ &= ~kEntriesSorted;
}

template <int kEmbeddingDim>
void InvertedFile<kEmbeddingDim>::SortEntries() {
  std::sort(entries_.begin(),
            entries_.end(),
            [](const EntryType& entry1, const EntryType& entry2) {
              return entry1.image_id < entry2.image_id;
            });
  status_ |= kEntriesSorted;
}

template <int kEmbeddingDim>
void InvertedFile<kEmbeddingDim>::ClearEntries() {
  entries_.clear();
  status_ &= ~kEntriesSorted;
}

template <int kEmbeddingDim>
void InvertedFile<kEmbeddingDim>::Reset() {
  status_ = kUnusable;
  squared_idf_weight_ = 0.0f;
  entries_.clear();
  thresholds_.setZero();
}

template <int kEmbeddingDim>
void InvertedFile<kEmbeddingDim>::ConvertToBinaryDescriptor(
    const DescType& descriptor,
    std::bitset<kEmbeddingDim>* binary_descriptor) const {
  THROW_CHECK_EQ(descriptor.size(), kEmbeddingDim);
  for (int i = 0; i < kEmbeddingDim; ++i) {
    (*binary_descriptor)[i] = descriptor[i] > thresholds_[i];
  }
}

template <int kEmbeddingDim>
void InvertedFile<kEmbeddingDim>::ComputeIDFWeight(const int num_total_images) {
  if (entries_.empty()) {
    squared_idf_weight_ = 0.0f;
    return;
  }

  std::unordered_set<int> image_ids;
  GetImageIds(&image_ids);

  const float idf_weight = std::log(static_cast<double>(num_total_images) /
                                    static_cast<double>(image_ids.size()));
  squared_idf_weight_ = idf_weight * idf_weight;
}

template <int kEmbeddingDim>
float InvertedFile<kEmbeddingDim>::SquaredIDFWeight() const {
  return squared_idf_weight_;
}

template <int kEmbeddingDim>
void InvertedFile<kEmbeddingDim>::ComputeHammingEmbedding(
    const Eigen::Matrix<float, Eigen::Dynamic, kEmbeddingDim>& descriptors) {
  const int num_descriptors = static_cast<int>(descriptors.rows());
  if (num_descriptors < 2) {
    return;
  }

  std::vector<float> elements(num_descriptors);
  for (int n = 0; n < kEmbeddingDim; ++n) {
    for (int i = 0; i < num_descriptors; ++i) {
      elements[i] = descriptors(i, n);
    }
    thresholds_[n] = Median(elements);
  }

  status_ |= kHasEmbedding;
}

template <int kEmbeddingDim>
void InvertedFile<kEmbeddingDim>::ScoreFeature(
    const DescType& descriptor, std::vector<ImageScore>* image_scores) const {
  THROW_CHECK_EQ(descriptor.size(), kEmbeddingDim);

  image_scores->clear();

  if (!IsUsable()) {
    return;
  }

  if (entries_.empty()) {
    return;
  }

  std::bitset<kEmbeddingDim> bin_descriptor;
  ConvertToBinaryDescriptor(descriptor, &bin_descriptor);

  ImageScore image_score;
  image_score.image_id = entries_.front().image_id;
  image_score.score = 0.0f;
  int num_image_votes = 0;

  // Finalizes the voting since we now know how many features from
  // the database image match the current image feature. This is
  // required to perform burstiness normalization (cf. Eqn. 2 in
  // Arandjelovic, Zisserman: Scalable descriptor
  // distinctiveness for location recognition. ACCV 2014).
  // Notice that the weight from the descriptor matching is already
  // accumulated in image_score.score, i.e., we only need
  // to apply the burstiness weighting.
  auto maybe_add_image_score =
      [this, &image_score, &image_scores, &num_image_votes]() {
        if (num_image_votes > 0) {
          image_score.score /= std::sqrt(static_cast<float>(num_image_votes));
          image_score.score *= squared_idf_weight_;
          image_scores->push_back(image_score);
        }
      };

  // Note that this assumes that the entries are sorted using SortEntries
  // according to their image identifiers.
  for (const auto& entry : entries_) {
    if (image_score.image_id < entry.image_id) {
      maybe_add_image_score();
      // Move to the next image.
      image_score.image_id = entry.image_id;
      image_score.score = 0.0f;
      num_image_votes = 0;
    }

    const size_t hamming_dist = (bin_descriptor ^ entry.descriptor).count();
    if (hamming_dist <= hamming_dist_weight_functor_.kMaxHammingDistance) {
      image_score.score += hamming_dist_weight_functor_(hamming_dist);
      num_image_votes += 1;
    }
  }

  // Add the voting for the largest image_id in the entries.
  maybe_add_image_score();
}

template <int kEmbeddingDim>
void InvertedFile<kEmbeddingDim>::GetImageIds(
    std::unordered_set<int>* ids) const {
  for (const EntryType& entry : entries_) {
    ids->insert(entry.image_id);
  }
}

template <int kEmbeddingDim>
void InvertedFile<kEmbeddingDim>::ComputeImageSelfSimilarities(
    std::unordered_map<int, double>* self_similarities) const {
  for (const auto& entry : entries_) {
    (*self_similarities)[entry.image_id] += squared_idf_weight_;
  }
}

template <int kEmbeddingDim>
void InvertedFile<kEmbeddingDim>::Read(std::istream* in) {
  THROW_CHECK(in->good());

  status_ = ReadBinaryLittleEndian<uint8_t>(in);

  float idf_weight = 0.0f;
  in->read(reinterpret_cast<char*>(&idf_weight), sizeof(float));
  squared_idf_weight_ = idf_weight * idf_weight;

  for (int i = 0; i < kEmbeddingDim; ++i) {
    in->read(reinterpret_cast<char*>(&thresholds_[i]), sizeof(float));
  }

  const uint32_t num_entries = ReadBinaryLittleEndian<uint32_t>(in);
  entries_.resize(num_entries);
  for (uint32_t i = 0; i < num_entries; ++i) {
    entries_[i].Read(in);
  }
}

template <int kEmbeddingDim>
void InvertedFile<kEmbeddingDim>::Write(std::ostream* out) const {
  THROW_CHECK(out->good());

  WriteBinaryLittleEndian<uint8_t>(out, status_);

  const float idf_weight = std::sqrt(squared_idf_weight_);
  out->write(reinterpret_cast<const char*>(&idf_weight), sizeof(float));

  for (int i = 0; i < kEmbeddingDim; ++i) {
    out->write(reinterpret_cast<const char*>(&thresholds_[i]), sizeof(float));
  }

  const uint32_t num_entries = entries_.size();
  WriteBinaryLittleEndian<uint32_t>(out, num_entries);
  for (uint32_t i = 0; i < num_entries; ++i) {
    entries_[i].Write(out);
  }
}

}  // namespace retrieval
}  // namespace colmap
