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

#include <climits>
#include <cstdio>
#include <memory>

#include <faiss/IVFlib.h>
#include <faiss/IndexAdditiveQuantizer.h>
#include <faiss/IndexIVFAdditiveQuantizer.h>
#include <faiss/MetricType.h>
#include <faiss/utils/distances.h>
#include <faiss/utils/hamming.h>
#include <faiss/utils/random.h>
#include <faiss/utils/utils.h>

/* This demo file shows how to:
 * - use a DistanceComputer to compute distances with encoded vectors
 * - in the context of an IVF, how to split an additive quantizer into an
 * AdditiveCoarseQuantizer and a ResidualQuantizer, in two different ways, with
 * and without storing the prefix.
 */

int main() {
    /******************************************
     * Generate a test dataset
     ******************************************/
    using idx_t = faiss::idx_t;
    size_t d = 128;
    size_t nt = 10000;
    size_t nb = 10000;
    size_t nq = 100;
    double t0 = faiss::getmillisecs();

    auto tic = [t0]() {
        printf("[%.3f s] ", (faiss::getmillisecs() - t0) / 1000);
    };

    tic();
    printf("sampling dataset of %zd dim vectors, Q %zd B %zd T %zd\n",
           d,
           nq,
           nb,
           nt);

    std::vector<float> buf(d * (nq + nt + nb));
    faiss::rand_smooth_vectors(nq + nt + nb, d, buf.data(), 1234);
    const float* xt = buf.data();
    const float* xb = buf.data() + nt * d;
    const float* xq = buf.data() + (nt + nb) * d;

    idx_t k = 10;
    std::vector<idx_t> gt(k * nq);
    std::vector<float> unused(k * nq);
    tic();
    printf("compute ground truth, k=%zd\n", k);
    faiss::knn_L2sqr(xq, xb, d, nq, nb, k, unused.data(), gt.data());

    // a function to compute the accuracy
    auto accuracy = [&](const idx_t* I) {
        idx_t accu = 0;
        for (idx_t q = 0; q < nq; q++) {
            accu += faiss::ranklist_intersection_size(
                    k, gt.data() + q * k, k, I + q * k);
        }
        return double(accu) / (k * nq);
    };

    /******************************************
     * Prepare the residual quantizer
     ******************************************/

    faiss::ResidualQuantizer rq(
            d, 7, 6, faiss::AdditiveQuantizer::ST_norm_qint8);
    // do cheap and inaccurate training
    rq.cp.niter = 5;
    rq.max_beam_size = 5;
    rq.train_type = 0;
    tic();
    printf("training the residual quantizer beam_size=%d\n", rq.max_beam_size);
    rq.train(nt, xt);

    tic();
    printf("encoding the database, code_size=%zd\n", rq.code_size);
    size_t code_size = rq.code_size;
    std::vector<uint8_t> raw_codes(nb * code_size);
    rq.compute_codes(xb, raw_codes.data(), nb);

    /****************************************************************
     * Make an index that uses that residual quantizer
     * Verify that a distance computer gives the same distances
     ****************************************************************/
    {
        faiss::IndexResidualQuantizer index(
                rq.d, rq.nbits, faiss::METRIC_L2, rq.search_type);

        // override trained index
        index.rq = rq;
        index.is_trained = true;

        // override vectors
        index.codes = faiss::MaybeOwnedVector<uint8_t>(raw_codes);
        index.ntotal = nb;

        tic();
        printf("IndexResidualQuantizer ready, searching\n");

        std::vector<float> D(k * nq);
        std::vector<idx_t> I(k * nq);
        index.search(nq, xq, k, D.data(), I.data());

        tic();
        printf("Accuracy (intersection @ %zd): %.3f\n", k, accuracy(I.data()));
        std::unique_ptr<faiss::FlatCodesDistanceComputer> dc(
                index.get_FlatCodesDistanceComputer());

        float max_diff12 = 0, max_diff13 = 0;

        for (idx_t q = 0; q < nq; q++) {
            const float* query = xq + q * d;
            dc->set_query(query);
            for (int i = 0; i < k; i++) {
                // 3 ways of computing the same distance

                // distance returned by the index
                float dis1 = D[q * k + i];

                // distance returned by the DistanceComputer that accesses the
                // index
                idx_t db_index = I[q * k + i];
                float dis2 = (*dc)(db_index);

                // distance computer from a code that does not belong to the
                // index
                const uint8_t* code = raw_codes.data() + code_size * db_index;
                float dis3 = dc->distance_to_code(code);

                max_diff12 = std::max(std::abs(dis1 - dis2), max_diff12);
                max_diff13 = std::max(std::abs(dis1 - dis3), max_diff13);
            }
        }
        tic();
        printf("Max DistanceComputer discrepancy 1-2: %g 1-3: %g\n",
               max_diff12,
               max_diff13);
    }

    /****************************************************************
     * Make an IVF index that uses the first 2 levels as a coarse quantizer
     * The IVF codes contain the full code (ie. redundant with the coarse
     *quantizer code)
     ****************************************************************/
    {
        // build a coarse quantizer from the 2 first levels of the RQ
        std::vector<size_t> nbits(2);
        std::copy(rq.nbits.begin(), rq.nbits.begin() + 2, nbits.begin());
        faiss::ResidualCoarseQuantizer rcq(rq.d, nbits);

        // set the coarse quantizer from the 2 first quantizers
        rcq.rq.initialize_from(rq);
        rcq.is_trained = true;
        rcq.ntotal = (idx_t)1 << rcq.rq.tot_bits;

        // settings for exhaustive search in RCQ
        rcq.centroid_norms.resize(rcq.ntotal);
        rcq.aq->compute_centroid_norms(rcq.centroid_norms.data());
        rcq.beam_factor = -1.0; // use exact search
        size_t nlist = rcq.ntotal;
        tic();
        printf("RCQ nlist = %zd tot_bits=%zd\n", nlist, rcq.rq.tot_bits);

        // build a IVFResidualQuantizer from that
        faiss::IndexIVFResidualQuantizer index(
                &rcq, rcq.d, nlist, rq.nbits, faiss::METRIC_L2, rq.search_type);
        index.by_residual = false;
        index.rq = rq;
        index.is_trained = true;

        // there are 3 ways of filling up the index...
        for (std::string filled_with : {"add", "manual", "derived"}) {
            tic();
            printf("filling up the index with %s, code_size=%zd\n",
                   filled_with.c_str(),
                   index.code_size);

            index.reset();

            if (filled_with == "add") {
                // standard add method
                index.add(nb, xb);
            } else if (filled_with == "manual") {
                // compute inverted lists and add elements manually
                // fill in the inverted index manually
                faiss::InvertedLists& invlists = *index.invlists;

                // assign vectors to inverted lists
                std::vector<idx_t> listnos(nb);
                std::vector<float> unused(nb);
                rcq.search(nb, xb, 1, unused.data(), listnos.data());

                // populate inverted lists
                for (idx_t i = 0; i < nb; i++) {
                    invlists.add_entry(
                            listnos[i], i, &raw_codes[i * code_size]);
                }

                index.ntotal = nb;
            } else if (filled_with == "derived") {
                // Since we have the raw codes precomputed, their prefix is the
                // inverted list index, so let's use that.
                faiss::InvertedLists& invlists = *index.invlists;

                // populate inverted lists
                for (idx_t i = 0; i < nb; i++) {
                    const uint8_t* code = &raw_codes[i * code_size];
                    faiss::BitstringReader rd(code, code_size);
                    idx_t list_no =
                            rd.read(rcq.rq.tot_bits); // read the list number
                    invlists.add_entry(list_no, i, code);
                }

                index.ntotal = nb;
            }

            tic();
            printf("Index filled in\n");

            for (int nprobe : {1, 4, 16, 64, int(nlist)}) {
                printf("setting nprobe=%-4d", nprobe);

                index.nprobe = nprobe;
                std::vector<float> D(k * nq);
                std::vector<idx_t> I(k * nq);
                index.search(nq, xq, k, D.data(), I.data());

                tic();
                printf("Accuracy (intersection @ %zd): %.3f\n",
                       k,
                       accuracy(I.data()));
            }
        }
    }

    /****************************************************************
     * Make an IVF index that uses the first 2 levels as a coarse
     * quantizer, but this time does not store the code prefix from the index
     ****************************************************************/

    {
        // build a coarse quantizer from the 2 first levels of the RQ
        int nlevel = 2;

        std::unique_ptr<faiss::IndexIVFResidualQuantizer> index(
                faiss::ivflib::ivf_residual_from_quantizer(rq, nlevel));

        // there are 2 ways of filling up the index...
        for (std::string filled_with : {"add", "derived"}) {
            tic();
            printf("filling up the IVF index with %s, code_size=%zd\n",
                   filled_with.c_str(),
                   index->code_size);

            index->reset();

            if (filled_with == "add") {
                // standard add method
                index->add(nb, xb);
            } else if (filled_with == "derived") {
                faiss::ivflib::ivf_residual_add_from_flat_codes(
                        index.get(), nb, raw_codes.data(), rq.code_size);
            }

            tic();
            printf("Index filled in\n");

            for (int nprobe : {1, 4, 16, 64, int(index->nlist)}) {
                printf("setting nprobe=%-4d", nprobe);

                index->nprobe = nprobe;
                std::vector<float> D(k * nq);
                std::vector<idx_t> I(k * nq);
                index->search(nq, xq, k, D.data(), I.data());

                tic();
                printf("Accuracy (intersection @ %zd): %.3f\n",
                       k,
                       accuracy(I.data()));
            }
        }
    }

    return 0;
}