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

#pragma once

#include <stdint.h>

#include <random>
#include <string>
#include <unordered_map>
#include <vector>

#include <faiss/impl/AdditiveQuantizer.h>
#include <faiss/impl/platform_macros.h>
#include <faiss/utils/utils.h>

namespace faiss {

namespace lsq {

struct IcmEncoderFactory;

} // namespace lsq

/** Implementation of LSQ/LSQ++ described in the following two papers:
 *
 * Revisiting additive quantization
 * Julieta Martinez, et al. ECCV 2016
 *
 * LSQ++: Lower running time and higher recall in multi-codebook quantization
 * Julieta Martinez, et al. ECCV 2018
 *
 * This implementation is mostly translated from the Julia implementations
 * by Julieta Martinez:
 * (https://github.com/una-dinosauria/local-search-quantization,
 *  https://github.com/una-dinosauria/Rayuela.jl)
 *
 * The trained codes are stored in `codebooks` which is called
 * `centroids` in PQ and RQ.
 */
struct LocalSearchQuantizer : AdditiveQuantizer {
    size_t K; ///< number of codes per codebook

    size_t train_iters = 25;      ///< number of iterations in training
    size_t encode_ils_iters = 16; ///< iterations of local search in encoding
    size_t train_ils_iters = 8;   ///< iterations of local search in training
    size_t icm_iters = 4;         ///< number of iterations in icm

    float p = 0.5f;      ///< temperature factor
    float lambd = 1e-2f; ///< regularization factor

    size_t chunk_size = 10000; ///< nb of vectors to encode at a time

    int random_seed = 0x12345; ///< seed for random generator
    size_t nperts = 4;         ///< number of perturbation in each code

    ///< if non-NULL, use this encoder to encode (owned by the object)
    lsq::IcmEncoderFactory* icm_encoder_factory = nullptr;

    bool update_codebooks_with_double = true;

    LocalSearchQuantizer(
            size_t d,     /* dimensionality of the input vectors */
            size_t M,     /* number of subquantizers */
            size_t nbits, /* number of bit per subvector index */
            Search_type_t search_type =
                    ST_decompress); /* determines the storage type */

    LocalSearchQuantizer();

    ~LocalSearchQuantizer() override;

    // Train the local search quantizer
    void train(size_t n, const float* x) override;

    /** Encode a set of vectors
     *
     * @param x      vectors to encode, size n * d
     * @param codes  output codes, size n * code_size
     * @param n      number of vectors
     * @param centroids  centroids to be added to x, size n * d
     */
    void compute_codes_add_centroids(
            const float* x,
            uint8_t* codes,
            size_t n,
            const float* centroids = nullptr) const override;

    /** Update codebooks given encodings
     *
     * @param x      training vectors, size n * d
     * @param codes  encoded training vectors, size n * M
     * @param n      number of vectors
     */
    void update_codebooks(const float* x, const int32_t* codes, size_t n);

    /** Encode vectors given codebooks using iterative conditional mode (icm).
     *
     * @param codes     output codes, size n * M
     * @param x         vectors to encode, size n * d
     * @param n         number of vectors
     * @param ils_iters number of iterations of iterative local search
     */
    void icm_encode(
            int32_t* codes,
            const float* x,
            size_t n,
            size_t ils_iters,
            std::mt19937& gen) const;

    void icm_encode_impl(
            int32_t* codes,
            const float* x,
            const float* unaries,
            std::mt19937& gen,
            size_t n,
            size_t ils_iters,
            bool verbose) const;

    void icm_encode_step(
            int32_t* codes,
            const float* unaries,
            const float* binaries,
            size_t n,
            size_t n_iters) const;

    /** Add some perturbation to codes
     *
     * @param codes codes to be perturbed, size n * M
     * @param n     number of vectors
     */
    void perturb_codes(int32_t* codes, size_t n, std::mt19937& gen) const;

    /** Add some perturbation to codebooks
     *
     * @param T         temperature of simulated annealing
     * @param stddev    standard deviations of each dimension in training data
     */
    void perturb_codebooks(
            float T,
            const std::vector<float>& stddev,
            std::mt19937& gen);

    /** Compute binary terms
     *
     * @param binaries binary terms, size M * M * K * K
     */
    void compute_binary_terms(float* binaries) const;

    /** Compute unary terms
     *
     * @param n       number of vectors
     * @param x       vectors to encode, size n * d
     * @param unaries unary terms, size n * M * K
     */
    void compute_unary_terms(const float* x, float* unaries, size_t n) const;

    /** Helper function to compute reconstruction error
     *
     * @param codes encoded codes, size n * M
     * @param x     vectors to encode, size n * d
     * @param n     number of vectors
     * @param objs  if it is not null, store reconstruction
                    error of each vector into it, size n
     */
    float evaluate(
            const int32_t* codes,
            const float* x,
            size_t n,
            float* objs = nullptr) const;
};

namespace lsq {

struct IcmEncoder {
    std::vector<float> binaries;

    bool verbose;

    const LocalSearchQuantizer* lsq;

    explicit IcmEncoder(const LocalSearchQuantizer* lsq);

    virtual ~IcmEncoder() {}

    ///< compute binary terms
    virtual void set_binary_term();

    /** Encode vectors given codebooks
     *
     * @param codes     output codes, size n * M
     * @param x         vectors to encode, size n * d
     * @param gen       random generator
     * @param n         number of vectors
     * @param ils_iters number of iterations of iterative local search
     */
    virtual void encode(
            int32_t* codes,
            const float* x,
            std::mt19937& gen,
            size_t n,
            size_t ils_iters) const;
};

struct IcmEncoderFactory {
    virtual IcmEncoder* get(const LocalSearchQuantizer* lsq) {
        return new IcmEncoder(lsq);
    }
    virtual ~IcmEncoderFactory() {}
};

/** A helper struct to count consuming time during training.
 *  It is NOT thread-safe.
 */
struct LSQTimer {
    std::unordered_map<std::string, double> t;

    LSQTimer() {
        reset();
    }

    double get(const std::string& name);

    void add(const std::string& name, double delta);

    void reset();
};

struct LSQTimerScope {
    double t0;
    LSQTimer* timer;
    std::string name;
    bool finished;

    LSQTimerScope(LSQTimer* timer, std::string name);

    void finish();

    ~LSQTimerScope();
};

} // namespace lsq

} // namespace faiss
