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

#ifndef FAISS_PRODUCT_QUANTIZER_H
#define FAISS_PRODUCT_QUANTIZER_H

#include <stdint.h>

#include <vector>

#include <faiss/Clustering.h>
#include <faiss/impl/Quantizer.h>
#include <faiss/impl/platform_macros.h>
#include <faiss/utils/Heap.h>

namespace faiss {

/** Product Quantizer.
 * PQ is trained using k-means, minimizing the L2 distance to centroids.
 * PQ supports L2 and Inner Product search, however the quantization error is
 * biased towards L2 distance.
 */
struct ProductQuantizer : Quantizer {
    size_t M;     ///< number of subquantizers
    size_t nbits; ///< number of bits per quantization index

    // values derived from the above
    size_t dsub;  ///< dimensionality of each subvector
    size_t ksub;  ///< number of centroids for each subquantizer
    bool verbose; ///< verbose during training?

    /// initialization
    enum train_type_t {
        Train_default,
        Train_hot_start,     ///< the centroids are already initialized
        Train_shared,        ///< share dictionary across PQ segments
        Train_hypercube,     ///< initialize centroids with nbits-D hypercube
        Train_hypercube_pca, ///< initialize centroids with nbits-D hypercube
    };
    train_type_t train_type;

    ClusteringParameters cp; ///< parameters used during clustering

    /// if non-NULL, use this index for assignment (should be of size
    /// d / M)
    Index* assign_index;

    /// Centroid table, size M * ksub * dsub.
    /// Layout: (M, ksub, dsub)
    std::vector<float> centroids;

    /// Transposed centroid table, size M * ksub * dsub.
    /// Layout: (dsub, M, ksub)
    std::vector<float> transposed_centroids;

    /// Squared lengths of centroids, size M * ksub
    /// Layout: (M, ksub)
    std::vector<float> centroids_sq_lengths;

    /// return the centroids associated with subvector m
    float* get_centroids(size_t m, size_t i) {
        return &centroids[(m * ksub + i) * dsub];
    }
    const float* get_centroids(size_t m, size_t i) const {
        return &centroids[(m * ksub + i) * dsub];
    }

    // Train the product quantizer on a set of points. A clustering
    // can be set on input to define non-default clustering parameters
    void train(size_t n, const float* x) override;

    ProductQuantizer(
            size_t d,      /* dimensionality of the input vectors */
            size_t M,      /* number of subquantizers */
            size_t nbits); /* number of bit per subvector index */

    ProductQuantizer();

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

    /// Define the centroids for subquantizer m
    void set_params(const float* centroids, int m);

    /// Quantize one vector with the product quantizer
    void compute_code(const float* x, uint8_t* code) const;

    /// same as compute_code for several vectors
    void compute_codes(const float* x, uint8_t* codes, size_t n) const override;

    /// speed up code assignment using assign_index
    /// (non-const because the index is changed)
    void compute_codes_with_assign_index(
            const float* x,
            uint8_t* codes,
            size_t n);

    /// decode a vector from a given code (or n vectors if third argument)
    void decode(const uint8_t* code, float* x) const;
    void decode(const uint8_t* code, float* x, size_t n) const override;

    /// If we happen to have the distance tables precomputed, this is
    /// more efficient to compute the codes.
    void compute_code_from_distance_table(const float* tab, uint8_t* code)
            const;

    /** Compute distance table for one vector.
     *
     * The distance table for x = [x_0 x_1 .. x_(M-1)] is a M * ksub
     * matrix that contains
     *
     *   dis_table (m, j) = || x_m - c_(m, j)||^2
     *   for m = 0..M-1 and j = 0 .. ksub - 1
     *
     * where c_(m, j) is the centroid no j of sub-quantizer m.
     *
     * @param x         input vector size d
     * @param dis_table output table, size M * ksub
     */
    void compute_distance_table(const float* x, float* dis_table) const;

    void compute_inner_prod_table(const float* x, float* dis_table) const;

    /** compute distance table for several vectors
     * @param nx        nb of input vectors
     * @param x         input vector size nx * d
     * @param dis_table output table, size nx * M * ksub
     */
    void compute_distance_tables(size_t nx, const float* x, float* dis_tables)
            const;

    void compute_inner_prod_tables(size_t nx, const float* x, float* dis_tables)
            const;

    /** perform a search (L2 distance)
     * @param x        query vectors, size nx * d
     * @param nx       nb of queries
     * @param codes    database codes, size ncodes * code_size
     * @param ncodes   nb of nb vectors
     * @param res      heap array to store results (nh == nx)
     * @param init_finalize_heap  initialize heap (input) and sort (output)?
     */
    void search(
            const float* x,
            size_t nx,
            const uint8_t* codes,
            const size_t ncodes,
            float_maxheap_array_t* res,
            bool init_finalize_heap = true) const;

    /** same as search, but with inner product similarity */
    void search_ip(
            const float* x,
            size_t nx,
            const uint8_t* codes,
            const size_t ncodes,
            float_minheap_array_t* res,
            bool init_finalize_heap = true) const;

    /// Symmetric Distance Table
    std::vector<float> sdc_table;

    // initialize the SDC table from the centroids
    void compute_sdc_table();

    void search_sdc(
            const uint8_t* qcodes,
            size_t nq,
            const uint8_t* bcodes,
            const size_t ncodes,
            float_maxheap_array_t* res,
            bool init_finalize_heap = true) const;

    /// Sync transposed centroids with regular centroids. This call
    /// is needed if centroids were edited directly.
    void sync_transposed_centroids();

    /// Clear transposed centroids table so ones are no longer used.
    void clear_transposed_centroids();
};

// block size used in ProductQuantizer::compute_codes
FAISS_API extern int product_quantizer_compute_codes_bs;

/*************************************************
 * Objects to encode / decode strings of bits
 *************************************************/

struct PQEncoderGeneric {
    uint8_t* code; ///< code for this vector
    uint8_t offset;
    const int nbits; ///< number of bits per subquantizer index

    uint8_t reg;

    PQEncoderGeneric(uint8_t* code, int nbits, uint8_t offset = 0);

    void encode(uint64_t x);

    ~PQEncoderGeneric();
};

struct PQEncoder8 {
    uint8_t* code;
    PQEncoder8(uint8_t* code, int nbits);
    void encode(uint64_t x);
};

struct PQEncoder16 {
    uint16_t* code;
    PQEncoder16(uint8_t* code, int nbits);
    void encode(uint64_t x);
};

struct PQDecoderGeneric {
    const uint8_t* code;
    uint8_t offset;
    const int nbits;
    const uint64_t mask;
    uint8_t reg;
    PQDecoderGeneric(const uint8_t* code, int nbits);
    uint64_t decode();
};

struct PQDecoder8 {
    static const int nbits = 8;
    const uint8_t* code;
    PQDecoder8(const uint8_t* code, int nbits);
    uint64_t decode();
};

struct PQDecoder16 {
    static const int nbits = 16;
    const uint16_t* code;
    PQDecoder16(const uint8_t* code, int nbits);
    uint64_t decode();
};

} // namespace faiss

#include <faiss/impl/ProductQuantizer-inl.h>

#endif
