/*
 * 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 <cmath>
#include <cstdint>
#include <vector>

#include <faiss/Index.h>
#include <faiss/IndexFlat.h>
#include <faiss/impl/Quantizer.h>

namespace faiss {

/** Abstract structure for additive quantizers
 *
 * Different from the product quantizer in which the decoded vector is the
 * concatenation of M sub-vectors, additive quantizers sum M sub-vectors
 * to get the decoded vector.
 */
struct AdditiveQuantizer : Quantizer {
    size_t M;                     ///< number of codebooks
    std::vector<size_t> nbits;    ///< bits for each step
    std::vector<float> codebooks; ///< codebooks

    // derived values
    /// codebook #1 is stored in rows codebook_offsets[i]:codebook_offsets[i+1]
    /// in the codebooks table of size total_codebook_size by d
    std::vector<uint64_t> codebook_offsets;
    size_t tot_bits = 0;            ///< total number of bits (indexes + norms)
    size_t norm_bits = 0;           ///< bits allocated for the norms
    size_t total_codebook_size = 0; ///< size of the codebook in vectors
    bool only_8bit = false;         ///< are all nbits = 8 (use faster decoder)

    bool verbose = false;    ///< verbose during training?
    bool is_trained = false; ///< is trained or not

    /// auxiliary data for ST_norm_lsq2x4 and ST_norm_rq2x4
    /// store norms of codebook entries for 4-bit fastscan
    std::vector<float> norm_tabs;
    IndexFlat1D qnorm; ///< store and search norms

    void compute_codebook_tables();

    /// norms of all codebook entries (size total_codebook_size)
    std::vector<float> centroid_norms;

    /// dot products of all codebook entries with the previous codebooks
    /// size sum(codebook_offsets[m] * 2^nbits[m], m=0..M-1)
    std::vector<float> codebook_cross_products;

    /// norms and distance matrixes with beam search can get large, so use this
    /// to control for the amount of memory that can be allocated
    size_t max_mem_distances = 5 * (size_t(1) << 30);

    /// encode a norm into norm_bits bits
    uint64_t encode_norm(float norm) const;

    /// encode norm by non-uniform scalar quantization
    uint32_t encode_qcint(float x) const;

    /// decode norm by non-uniform scalar quantization
    float decode_qcint(uint32_t c) const;

    /// Encodes how search is performed and how vectors are encoded
    enum Search_type_t {
        ST_decompress,    ///< decompress database vector
        ST_LUT_nonorm,    ///< use a LUT, don't include norms (OK for IP or
                          ///< normalized vectors)
        ST_norm_from_LUT, ///< compute the norms from the look-up tables (cost
                          ///< is in O(M^2))
        ST_norm_float, ///< use a LUT, and store float32 norm with the vectors
        ST_norm_qint8, ///< use a LUT, and store 8bit-quantized norm
        ST_norm_qint4,
        ST_norm_cqint8, ///< use a LUT, and store non-uniform quantized norm
        ST_norm_cqint4,

        ST_norm_lsq2x4, ///< use a 2x4 bits lsq as norm quantizer (for fast
                        ///< scan)
        ST_norm_rq2x4,  ///< use a 2x4 bits rq as norm quantizer (for fast scan)
    };

    AdditiveQuantizer(
            size_t d,
            const std::vector<size_t>& nbits,
            Search_type_t search_type = ST_decompress);

    AdditiveQuantizer();

    ///< compute derived values when d, M and nbits have been set
    void set_derived_values();

    ///< Train the norm quantizer
    void train_norm(size_t n, const float* norms);

    void compute_codes(const float* x, uint8_t* codes, size_t n)
            const override {
        compute_codes_add_centroids(x, codes, n);
    }

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

    /** pack a series of code to bit-compact format
     *
     * @param codes        codes to be packed, size n * code_size
     * @param packed_codes output bit-compact codes
     * @param ld_codes     leading dimension of codes
     * @param norms        norms of the vectors (size n). Will be computed if
     *                     needed but not provided
     * @param centroids    centroids to be added to x, size n * d
     */
    void pack_codes(
            size_t n,
            const int32_t* codes,
            uint8_t* packed_codes,
            int64_t ld_codes = -1,
            const float* norms = nullptr,
            const float* centroids = nullptr) const;

    /** Decode a set of vectors
     *
     * @param codes  codes to decode, size n * code_size
     * @param x      output vectors, size n * d
     */
    void decode(const uint8_t* codes, float* x, size_t n) const override;

    /** Decode a set of vectors in non-packed format
     *
     * @param codes  codes to decode, size n * ld_codes
     * @param x      output vectors, size n * d
     */
    virtual void decode_unpacked(
            const int32_t* codes,
            float* x,
            size_t n,
            int64_t ld_codes = -1) const;

    /****************************************************************************
     * Search functions in an external set of codes.
     ****************************************************************************/

    /// Also determines what's in the codes
    Search_type_t search_type;

    /// min/max for quantization of norms
    float norm_min = NAN, norm_max = NAN;

    template <bool is_IP, Search_type_t effective_search_type>
    float compute_1_distance_LUT(const uint8_t* codes, const float* LUT) const;

    /*
        float compute_1_L2sqr(const uint8_t* codes, const float* LUT);
    */
    /****************************************************************************
     * Support for exhaustive distance computations with all the centroids.
     * Hence, the number of these centroids should not be too large.
     ****************************************************************************/

    /// decoding function for a code in a 64-bit word
    void decode_64bit(idx_t n, float* x) const;

    /** Compute inner-product look-up tables. Used in the centroid search
     * functions.
     *
     * @param xq     query vector, size (n, d)
     * @param LUT    look-up table, size (n, total_codebook_size)
     * @param alpha  compute alpha * inner-product
     * @param ld_lut  leading dimension of LUT
     */
    virtual void compute_LUT(
            size_t n,
            const float* xq,
            float* LUT,
            float alpha = 1.0f,
            long ld_lut = -1) const;

    /// exact IP search
    void knn_centroids_inner_product(
            idx_t n,
            const float* xq,
            idx_t k,
            float* distances,
            idx_t* labels) const;

    /** For L2 search we need the L2 norms of the centroids
     *
     * @param norms    output norms table, size total_codebook_size
     */
    void compute_centroid_norms(float* norms) const;

    /** Exact L2 search, with precomputed norms */
    void knn_centroids_L2(
            idx_t n,
            const float* xq,
            idx_t k,
            float* distances,
            idx_t* labels,
            const float* centroid_norms) const;

    virtual ~AdditiveQuantizer() override;
};

} // namespace faiss
