/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This source code is licensed under the MIT license found in the
 * LICENSE file in the root directory of this source tree.
 */

// -*- c++ -*-

#pragma once

#include <algorithm>
#include <mutex>
#include <queue>
#include <random>
#include <unordered_set>
#include <vector>

#include <omp.h>

#include <faiss/Index.h>
#include <faiss/impl/FaissAssert.h>
#include <faiss/impl/platform_macros.h>
#include <faiss/utils/Heap.h>
#include <faiss/utils/random.h>

namespace faiss {

/** Implementation of NNDescent which is one of the most popular
 *  KNN graph building algorithms
 *
 * Efficient K-Nearest Neighbor Graph Construction for Generic
 * Similarity Measures
 *
 *  Dong, Wei, Charikar Moses, and Kai Li, WWW 2011
 *
 * This implementation is heavily influenced by the efanna
 * implementation by Cong Fu and the KGraph library by Wei Dong
 * (https://github.com/ZJULearning/efanna_graph)
 * (https://github.com/aaalgo/kgraph)
 *
 * The NNDescent object stores only the neighbor link structure,
 * see IndexNNDescent.h for the full index object.
 */

struct VisitedTable;
struct DistanceComputer;

namespace nndescent {

struct Neighbor {
    int id;
    float distance;
    bool flag;

    Neighbor() = default;
    Neighbor(int id, float distance, bool f)
            : id(id), distance(distance), flag(f) {}

    inline bool operator<(const Neighbor& other) const {
        return distance < other.distance;
    }
};

struct Nhood {
    std::mutex lock;
    std::vector<Neighbor> pool; // candidate pool (a max heap)
    int M;                      // number of new neighbors to be operated

    std::vector<int> nn_old;  // old neighbors
    std::vector<int> nn_new;  // new neighbors
    std::vector<int> rnn_old; // reverse old neighbors
    std::vector<int> rnn_new; // reverse new neighbors

    Nhood() = default;

    Nhood(int l, int s, std::mt19937& rng, int N);

    Nhood& operator=(const Nhood& other);

    Nhood(const Nhood& other);

    void insert(int id, float dist);

    template <typename C>
    void join(C callback) const;
};

} // namespace nndescent

struct NNDescent {
    using storage_idx_t = int;

    using KNNGraph = std::vector<nndescent::Nhood>;

    explicit NNDescent(const int d, const int K);

    ~NNDescent();

    void build(DistanceComputer& qdis, const int n, bool verbose);

    void search(
            DistanceComputer& qdis,
            const int topk,
            idx_t* indices,
            float* dists,
            VisitedTable& vt) const;

    void reset();

    /// Initialize the KNN graph randomly
    void init_graph(DistanceComputer& qdis);

    /// Perform NNDescent algorithm
    void nndescent(DistanceComputer& qdis, bool verbose);

    /// Perform local join on each node
    void join(DistanceComputer& qdis);

    /// Sample new neighbors for each node to perform local join later
    void update();

    /// Sample a small number of points to evaluate the quality of KNNG built
    void generate_eval_set(
            DistanceComputer& qdis,
            std::vector<int>& c,
            std::vector<std::vector<int>>& v,
            int N);

    /// Evaluate the quality of KNNG built
    float eval_recall(
            std::vector<int>& ctrl_points,
            std::vector<std::vector<int>>& acc_eval_set);

    bool has_built = false;

    int S = 10;  // number of sample neighbors to be updated for each node
    int R = 100; // size of reverse links, 0 means the reverse links will not be
                 // used
    int iter = 10;          // number of iterations to iterate over
    int search_L = 0;       // size of candidate pool in searching
    int random_seed = 2021; // random seed for generators

    int K; // K in KNN graph
    int d; // dimensions
    int L; // size of the candidate pool in building

    int ntotal = 0;

    KNNGraph graph;
    std::vector<int> final_graph;
};

} // namespace faiss
