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

#include <algorithm>
#include <array>
#include <cmath>
#include <cstring>
#include <memory>
#include <random>
#include <vector>

#include <faiss/IndexFlat.h>
#include <faiss/IndexIVFFlat.h>
#include <faiss/IndexScalarQuantizer.h>
#include <faiss/impl/ScalarQuantizer.h>
#include <faiss/impl/scalar_quantizer/distance_computers.h>
#include <faiss/impl/scalar_quantizer/quantizers.h>
#include <faiss/impl/scalar_quantizer/scanners.h>
#include <faiss/impl/scalar_quantizer/similarities.h>
#include <faiss/index_factory.h>
#include <faiss/utils/distances.h>
#include <faiss/utils/random.h>
#include <faiss/utils/simd_levels.h>

namespace {

std::vector<float> make_normalized_vectors(size_t n, size_t d) {
    std::vector<float> x(n * d);
    faiss::float_randn(x.data(), x.size(), 1234);
    faiss::fvec_renorm_L2(d, n, x.data());
    return x;
}

float mean_squared_error(
        const std::vector<float>& x,
        const std::vector<float>& y) {
    EXPECT_EQ(x.size(), y.size());
    return faiss::fvec_L2sqr(x.data(), y.data(), x.size()) / x.size();
}

template <int NBits>
using ScalarLloydMaxQuantizer = faiss::scalar_quantizer::
        QuantizerLloydMax<NBits, faiss::SIMDLevel::NONE>;

template <int NBits>
using ScalarLloydMaxL2DistanceComputer = faiss::scalar_quantizer::DCTemplate<
        ScalarLloydMaxQuantizer<NBits>,
        faiss::scalar_quantizer::SimilarityL2<faiss::SIMDLevel::NONE>,
        faiss::SIMDLevel::NONE>;

template <int NBits>
using ScalarLloydMaxL2Scanner = faiss::scalar_quantizer::IVFSQScannerL2<
        ScalarLloydMaxL2DistanceComputer<NBits>>;

faiss::ScalarQuantizer make_trained_lloyd_max_sq(
        size_t d,
        faiss::ScalarQuantizer::QuantizerType qtype) {
    const size_t n = 128;
    std::vector<float> x = make_normalized_vectors(n, d);
    faiss::ScalarQuantizer sq(d, qtype);
    sq.train(n, x.data());
    return sq;
}

struct ScopedSIMDLevel {
    faiss::SIMDLevel original_level;

    explicit ScopedSIMDLevel(faiss::SIMDLevel level)
            : original_level(faiss::SIMDConfig::get_level()) {
        faiss::SIMDConfig::set_level(level);
    }

    ~ScopedSIMDLevel() {
        faiss::SIMDConfig::set_level(original_level);
    }
};

std::vector<faiss::SIMDLevel> available_lloyd_max_simd_levels() {
    std::vector<faiss::SIMDLevel> levels;
    for (faiss::SIMDLevel level :
         {faiss::SIMDLevel::AVX512,
          faiss::SIMDLevel::AVX2,
          faiss::SIMDLevel::ARM_NEON}) {
        if (faiss::SIMDConfig::is_simd_level_available(level)) {
            levels.push_back(level);
        }
    }
    return levels;
}

// SIMD levels with scalar-quantizer kernels that can be checked for parity
// against the scalar (NONE) path. Unlike dispatch tests, parity tests do not
// require dimension restrictions, so RISCV_RVV (no alignment constraints)
// is included here.
std::vector<faiss::SIMDLevel> available_sq_parity_simd_levels() {
    std::vector<faiss::SIMDLevel> levels;
    for (faiss::SIMDLevel level :
         {faiss::SIMDLevel::AVX512,
          faiss::SIMDLevel::AVX2,
          faiss::SIMDLevel::ARM_NEON,
          faiss::SIMDLevel::RISCV_RVV}) {
        if (faiss::SIMDConfig::is_simd_level_available(level)) {
            levels.push_back(level);
        }
    }
    return levels;
}

template <int NBits>
void expect_lloyd_max_simd_dispatch_for_compatible_dim(
        faiss::SIMDLevel level,
        faiss::ScalarQuantizer::QuantizerType qtype) {
    ScopedSIMDLevel scoped(level);
    const size_t d = 32;
    faiss::ScalarQuantizer sq = make_trained_lloyd_max_sq(d, qtype);

    std::unique_ptr<faiss::ScalarQuantizer::SQuantizer> quantizer(
            sq.select_quantizer());
    ASSERT_NE(quantizer, nullptr);

    std::unique_ptr<faiss::ScalarQuantizer::SQDistanceComputer> dc(
            sq.get_distance_computer(faiss::METRIC_L2));
    ASSERT_NE(dc, nullptr);

    std::unique_ptr<faiss::InvertedListScanner> scanner(
            sq.select_InvertedListScanner(
                    faiss::METRIC_L2, nullptr, false, nullptr, false));
    ASSERT_NE(scanner, nullptr);

    auto* quantizer_raw = quantizer.get();
    auto* dc_raw = dc.get();
    auto* scanner_raw = scanner.get();

    EXPECT_NE(typeid(*quantizer_raw), typeid(ScalarLloydMaxQuantizer<NBits>));
    EXPECT_NE(typeid(*dc_raw), typeid(ScalarLloydMaxL2DistanceComputer<NBits>));
    EXPECT_NE(typeid(*scanner_raw), typeid(ScalarLloydMaxL2Scanner<NBits>));
}

template <int NBits>
void expect_lloyd_max_simd_dispatch_fallback_for_incompatible_dim(
        faiss::SIMDLevel level,
        size_t d,
        faiss::ScalarQuantizer::QuantizerType qtype) {
    ScopedSIMDLevel scoped(level);
    faiss::ScalarQuantizer sq = make_trained_lloyd_max_sq(d, qtype);

    std::unique_ptr<faiss::ScalarQuantizer::SQuantizer> quantizer(
            sq.select_quantizer());
    ASSERT_NE(quantizer, nullptr);

    std::unique_ptr<faiss::ScalarQuantizer::SQDistanceComputer> dc(
            sq.get_distance_computer(faiss::METRIC_L2));
    ASSERT_NE(dc, nullptr);

    std::unique_ptr<faiss::InvertedListScanner> scanner(
            sq.select_InvertedListScanner(
                    faiss::METRIC_L2, nullptr, false, nullptr, false));
    ASSERT_NE(scanner, nullptr);

    auto* quantizer_raw = quantizer.get();
    auto* dc_raw = dc.get();
    auto* scanner_raw = scanner.get();

    EXPECT_EQ(typeid(*quantizer_raw), typeid(ScalarLloydMaxQuantizer<NBits>));
    EXPECT_EQ(typeid(*dc_raw), typeid(ScalarLloydMaxL2DistanceComputer<NBits>));
    EXPECT_EQ(typeid(*scanner_raw), typeid(ScalarLloydMaxL2Scanner<NBits>));
}

template <int NBits>
void check_lloyd_max_distance_path_parity(
        faiss::SIMDLevel level,
        faiss::ScalarQuantizer::QuantizerType qtype) {
    ScopedSIMDLevel scoped(level);
    const size_t d = 32;
    const size_t n = 128;
    std::vector<float> xb = make_normalized_vectors(n, d);
    std::vector<float> xq = make_normalized_vectors(1, d);

    faiss::ScalarQuantizer sq(d, qtype);
    sq.train(n, xb.data());

    std::vector<uint8_t> codes(sq.code_size * n, 0);
    sq.compute_codes(xb.data(), codes.data(), n);

    std::unique_ptr<faiss::ScalarQuantizer::SQDistanceComputer> scalar_dc(
            faiss::scalar_quantizer::sq_select_distance_computer<
                    faiss::SIMDLevel::NONE>(
                    faiss::METRIC_L2, qtype, d, sq.trained));
    std::unique_ptr<faiss::ScalarQuantizer::SQDistanceComputer> simd_dc(
            sq.get_distance_computer(faiss::METRIC_L2));
    ASSERT_NE(scalar_dc, nullptr);
    ASSERT_NE(simd_dc, nullptr);

    scalar_dc->set_query(xq.data());
    simd_dc->set_query(xq.data());

    std::vector<uint8_t> zero_code(sq.code_size, 0);
    std::vector<uint8_t> max_code(sq.code_size, 0xff);
    const std::array<const uint8_t*, 6> query_codes = {
            codes.data(),
            codes.data() + sq.code_size,
            codes.data() + 2 * sq.code_size,
            codes.data() + 3 * sq.code_size,
            zero_code.data(),
            max_code.data()};

    for (const uint8_t* code : query_codes) {
        EXPECT_NEAR(
                scalar_dc->query_to_code(code),
                simd_dc->query_to_code(code),
                1e-3);
    }

    std::vector<uint8_t> bundle(4 * sq.code_size, 0);
    std::memcpy(bundle.data(), zero_code.data(), sq.code_size);
    std::memcpy(bundle.data() + sq.code_size, max_code.data(), sq.code_size);
    std::memcpy(bundle.data() + 2 * sq.code_size, codes.data(), sq.code_size);
    std::memcpy(
            bundle.data() + 3 * sq.code_size,
            codes.data() + sq.code_size,
            sq.code_size);

    scalar_dc->codes = bundle.data();
    scalar_dc->code_size = sq.code_size;
    simd_dc->codes = bundle.data();
    simd_dc->code_size = sq.code_size;

    const std::array<std::pair<faiss::idx_t, faiss::idx_t>, 4> pairs = {{
            {0, 1},
            {0, 2},
            {1, 3},
            {2, 3},
    }};
    for (const auto& [lhs, rhs] : pairs) {
        EXPECT_NEAR(
                scalar_dc->symmetric_dis(lhs, rhs),
                simd_dc->symmetric_dis(lhs, rhs),
                1e-3);
    }

    std::unique_ptr<faiss::InvertedListScanner> scalar_scanner(
            faiss::scalar_quantizer::sq_select_InvertedListScanner<
                    faiss::SIMDLevel::NONE>(
                    qtype,
                    faiss::METRIC_L2,
                    d,
                    sq.code_size,
                    sq.trained,
                    nullptr,
                    false,
                    nullptr,
                    false));
    std::unique_ptr<faiss::InvertedListScanner> simd_scanner(
            sq.select_InvertedListScanner(
                    faiss::METRIC_L2, nullptr, false, nullptr, false));
    ASSERT_NE(scalar_scanner, nullptr);
    ASSERT_NE(simd_scanner, nullptr);

    scalar_scanner->set_query(xq.data());
    simd_scanner->set_query(xq.data());
    scalar_scanner->set_list(0, 0.0f);
    simd_scanner->set_list(0, 0.0f);

    for (const uint8_t* code : query_codes) {
        EXPECT_NEAR(
                scalar_scanner->distance_to_code(code),
                simd_scanner->distance_to_code(code),
                1e-3);
    }
}

// Parity between a SIMD level and the scalar (NONE) path for one quantizer
// type, metric and dimension. Covers quantizer decoding, the
// SQDistanceComputer APIs (query_to_code, query_to_codes_batch_4,
// symmetric_dis) and the inverted-list scanner's distance_to_code.
void check_sq_distance_path_parity(
        faiss::SIMDLevel level,
        faiss::ScalarQuantizer::QuantizerType qtype,
        faiss::MetricType metric,
        size_t d) {
    ScopedSIMDLevel scoped(level);
    const size_t n = 64;
    std::vector<float> xb = make_normalized_vectors(n, d);
    std::vector<float> xq = make_normalized_vectors(1, d);

    faiss::ScalarQuantizer sq(d, qtype);
    sq.train(n, xb.data());

    std::vector<uint8_t> codes(sq.code_size * n, 0);
    sq.compute_codes(xb.data(), codes.data(), n);

    // fp16/bf16 codes are raw IEEE bit patterns: arbitrary bytes can decode
    // to NaN, so only codes produced by compute_codes are checked for those.
    const bool arbitrary_bytes_ok = qtype != faiss::ScalarQuantizer::QT_fp16 &&
            qtype != faiss::ScalarQuantizer::QT_bf16;

    std::vector<uint8_t> zero_code(sq.code_size, 0);
    std::vector<uint8_t> max_code(sq.code_size, 0xff);
    std::vector<const uint8_t*> query_codes = {
            codes.data(),
            codes.data() + sq.code_size,
            codes.data() + 2 * sq.code_size,
            codes.data() + 3 * sq.code_size};
    if (arbitrary_bytes_ok) {
        query_codes.push_back(zero_code.data());
        query_codes.push_back(max_code.data());
    }

    std::unique_ptr<faiss::ScalarQuantizer::SQuantizer> scalar_quant(
            faiss::scalar_quantizer::sq_select_quantizer<
                    faiss::SIMDLevel::NONE>(qtype, d, sq.trained));
    std::unique_ptr<faiss::ScalarQuantizer::SQuantizer> simd_quant(
            sq.select_quantizer());
    ASSERT_NE(scalar_quant, nullptr);
    ASSERT_NE(simd_quant, nullptr);

    for (const uint8_t* code : query_codes) {
        std::vector<float> ref(d), out(d);
        scalar_quant->decode_vector(code, ref.data());
        simd_quant->decode_vector(code, out.data());
        for (size_t j = 0; j < d; j++) {
            EXPECT_NEAR(
                    ref[j], out[j], 1e-5 * std::max(1.0f, std::fabs(ref[j])));
        }
    }

    auto expect_distance_near = [](float ref, float out) {
        EXPECT_NEAR(ref, out, 5e-3 * std::max(1.0f, std::fabs(ref)));
    };

    std::unique_ptr<faiss::ScalarQuantizer::SQDistanceComputer> scalar_dc(
            faiss::scalar_quantizer::sq_select_distance_computer<
                    faiss::SIMDLevel::NONE>(metric, qtype, d, sq.trained));
    std::unique_ptr<faiss::ScalarQuantizer::SQDistanceComputer> simd_dc(
            sq.get_distance_computer(metric));
    ASSERT_NE(scalar_dc, nullptr);
    ASSERT_NE(simd_dc, nullptr);

    scalar_dc->set_query(xq.data());
    simd_dc->set_query(xq.data());

    for (const uint8_t* code : query_codes) {
        expect_distance_near(
                scalar_dc->query_to_code(code), simd_dc->query_to_code(code));
    }

    float scalar_dis[4], simd_dis[4];
    scalar_dc->query_to_codes_batch_4(
            query_codes[0],
            query_codes[1],
            query_codes[2],
            query_codes[3],
            scalar_dis[0],
            scalar_dis[1],
            scalar_dis[2],
            scalar_dis[3]);
    simd_dc->query_to_codes_batch_4(
            query_codes[0],
            query_codes[1],
            query_codes[2],
            query_codes[3],
            simd_dis[0],
            simd_dis[1],
            simd_dis[2],
            simd_dis[3]);
    for (int k = 0; k < 4; k++) {
        expect_distance_near(scalar_dis[k], simd_dis[k]);
    }

    std::vector<uint8_t> bundle(4 * sq.code_size, 0);
    for (int k = 0; k < 4; k++) {
        std::memcpy(
                bundle.data() + k * sq.code_size, query_codes[k], sq.code_size);
    }
    scalar_dc->codes = bundle.data();
    scalar_dc->code_size = sq.code_size;
    simd_dc->codes = bundle.data();
    simd_dc->code_size = sq.code_size;

    const std::array<std::pair<faiss::idx_t, faiss::idx_t>, 4> pairs = {{
            {0, 1},
            {0, 2},
            {1, 3},
            {2, 3},
    }};
    for (const auto& [lhs, rhs] : pairs) {
        expect_distance_near(
                scalar_dc->symmetric_dis(lhs, rhs),
                simd_dc->symmetric_dis(lhs, rhs));
    }

    std::unique_ptr<faiss::InvertedListScanner> scalar_scanner(
            faiss::scalar_quantizer::sq_select_InvertedListScanner<
                    faiss::SIMDLevel::NONE>(
                    qtype,
                    metric,
                    d,
                    sq.code_size,
                    sq.trained,
                    nullptr,
                    false,
                    nullptr,
                    false));
    std::unique_ptr<faiss::InvertedListScanner> simd_scanner(
            sq.select_InvertedListScanner(
                    metric, nullptr, false, nullptr, false));
    ASSERT_NE(scalar_scanner, nullptr);
    ASSERT_NE(simd_scanner, nullptr);

    scalar_scanner->set_query(xq.data());
    simd_scanner->set_query(xq.data());
    scalar_scanner->set_list(0, 0.0f);
    simd_scanner->set_list(0, 0.0f);

    for (const uint8_t* code : query_codes) {
        expect_distance_near(
                scalar_scanner->distance_to_code(code),
                simd_scanner->distance_to_code(code));
    }
}

void check_tqmse_roundtrip(
        size_t d,
        faiss::ScalarQuantizer::QuantizerType qtype) {
    const size_t n = 128;
    std::vector<float> x = make_normalized_vectors(n, d);
    faiss::ScalarQuantizer sq(d, qtype);

    sq.train(n, x.data());

    std::vector<uint8_t> codes(sq.code_size * n, 0);
    sq.compute_codes(x.data(), codes.data(), n);

    std::vector<float> decoded(n * d);
    sq.decode(codes.data(), decoded.data(), n);

    for (float v : decoded) {
        EXPECT_TRUE(std::isfinite(v));
        EXPECT_LE(v, 1.0f);
        EXPECT_GE(v, -1.0f);
    }
}

void check_eden_roundtrip(
        size_t d,
        size_t nb_bits,
        faiss::ScalarQuantizer::QuantizerType qtype) {
    const size_t n = 128;
    std::vector<float> x = make_normalized_vectors(n, d);
    faiss::ScalarQuantizer sq(d, qtype);

    sq.train(0, nullptr);

    const size_t k = size_t(1) << nb_bits;
    EXPECT_EQ(sq.bits, nb_bits);
    EXPECT_EQ(sq.trained.size(), 2 * k - 1);

    std::vector<uint8_t> codes(sq.code_size * n, 0);
    sq.compute_codes(x.data(), codes.data(), n);

    std::vector<float> decoded(n * d);
    sq.decode(codes.data(), decoded.data(), n);

    for (float v : decoded) {
        EXPECT_TRUE(std::isfinite(v));
    }
}

} // namespace

TEST(ScalarQuantizer, RSQuantilesClamping) {
    int d = 8;
    int n = 100;

    std::vector<float> x(d * n);
    for (size_t i = 0; i < x.size(); i++) {
        x[i] = static_cast<float>(i % 100);
    }

    faiss::ScalarQuantizer sq(d, faiss::ScalarQuantizer::QT_8bit);
    sq.rangestat = faiss::ScalarQuantizer::RS_quantiles;

    sq.rangestat_arg = 0.05f;
    ASSERT_NO_THROW(sq.train(n, x.data()));

    sq.rangestat_arg = 0.0f;
    ASSERT_NO_THROW(sq.train(n, x.data()));

    sq.rangestat_arg = -0.1f;
    ASSERT_NO_THROW(sq.train(n, x.data()));

    sq.rangestat_arg = 0.8f;
    ASSERT_NO_THROW(sq.train(n, x.data()));

    sq.rangestat_arg = 0.5f;
    ASSERT_NO_THROW(sq.train(n, x.data()));
}

TEST(ScalarQuantizer, RSQuantilesOddSize) {
    int d = 4;
    int n = 5;

    std::vector<float> x(d * n);
    for (size_t i = 0; i < x.size(); i++) {
        x[i] = static_cast<float>(i);
    }

    faiss::ScalarQuantizer sq(d, faiss::ScalarQuantizer::QT_8bit);
    sq.rangestat = faiss::ScalarQuantizer::RS_quantiles;

    sq.rangestat_arg = 0.4f;
    ASSERT_NO_THROW(sq.train(n, x.data()));

    sq.rangestat_arg = 0.5f;
    ASSERT_NO_THROW(sq.train(n, x.data()));

    sq.rangestat_arg = 0.6f;
    ASSERT_NO_THROW(sq.train(n, x.data()));
}

TEST(ScalarQuantizer, RSQuantilesValidRange) {
    int d = 8;
    int n = 100;

    std::vector<float> x(d * n);
    for (size_t i = 0; i < x.size(); i++) {
        x[i] = static_cast<float>(i);
    }

    faiss::ScalarQuantizer sq(d, faiss::ScalarQuantizer::QT_8bit);
    sq.rangestat = faiss::ScalarQuantizer::RS_quantiles;
    sq.rangestat_arg = 0.1f;

    sq.train(n, x.data());

    std::vector<uint8_t> codes(sq.code_size * n);
    ASSERT_NO_THROW(sq.compute_codes(x.data(), codes.data(), n));

    std::vector<float> decoded(d * n);
    ASSERT_NO_THROW(sq.decode(codes.data(), decoded.data(), n));
}

TEST(ScalarQuantizer, RSQuantilesSmallDataset) {
    int d = 2;
    int n = 2;

    std::vector<float> x = {1.0f, 2.0f, 3.0f, 4.0f};

    faiss::ScalarQuantizer sq(d, faiss::ScalarQuantizer::QT_8bit);
    sq.rangestat = faiss::ScalarQuantizer::RS_quantiles;
    sq.rangestat_arg = 0.1f;

    ASSERT_NO_THROW(sq.train(n, x.data()));
}

TEST(ScalarQuantizer, MinimalTraining) {
    std::vector<float> x = {1.0f, 2.0f, 3.0f, 4.0f};
    using SQ = faiss::ScalarQuantizer;
    for (auto qt : {
                 SQ::QT_4bit,
                 SQ::QT_4bit_uniform,
                 SQ::QT_6bit,
                 SQ::QT_8bit,
                 SQ::QT_8bit_uniform,
         }) {
        SCOPED_TRACE(qt);
        for (size_t d : {1, 2}) {
            faiss::IndexScalarQuantizer index(d, qt);
            index.sq.rangestat = SQ::RS_minmax;
            // n=0 is insufficient.
            EXPECT_THROW(index.train(0, x.data()), faiss::FaissException);
            // nullptr throws, not crashes.
            EXPECT_THROW(index.train(1, nullptr), faiss::FaissException);
            // 1x1 values is allowed.
            EXPECT_NO_THROW(index.train(1, x.data()));
        }
    }
}

TEST(ScalarQuantizer, RVVDirect8bitEncodingMatchesTruncation) {
    if (!faiss::SIMDConfig::is_simd_level_available(
                faiss::SIMDLevel::RISCV_RVV)) {
        GTEST_SKIP() << "RVV is not available";
    }
    ScopedSIMDLevel scoped(faiss::SIMDLevel::RISCV_RVV);
    // Negative fractions greater than -1 truncate to the valid code zero.
    // Inputs outside (-1, 256) have no portable scalar conversion contract.
    const std::vector<float> values = {
            std::nextafter(-1.0f, 0.0f),
            -0.75f,
            -0.0f,
            std::nextafter(0.0f, -1.0f),
            0.0f,
            std::nextafter(0.0f, 1.0f),
            std::nextafter(1.0f, 0.0f),
            1.0f,
            std::nextafter(1.0f, 2.0f),
            127.75f,
            128.5f,
            254.999f,
            255.0f,
            std::nextafter(256.0f, 0.0f)};
    for (size_t d : {0, 1, 15, 16, 17, 31, 32, 33, 65, 769}) {
        SCOPED_TRACE(d);
        const size_t n = 3;
        std::vector<float> input(d * n + 1);
        std::vector<uint8_t> expected(d * n + 32, 0xa5);
        for (size_t i = 0; i < d * n; ++i) {
            input[i] = values[i % values.size()];
            expected[16 + i] = static_cast<uint8_t>(input[i]);
        }
        std::vector<uint8_t> actual(d * n + 32, 0xa5);
        faiss::ScalarQuantizer sq(d, faiss::ScalarQuantizer::QT_8bit_direct);
        sq.compute_codes(input.data(), actual.data() + 16, n);
        EXPECT_EQ(actual, expected);
    }
}

TEST(TestSQ0bit, CoarseOnlySearch) {
    // Test QT_0bit: centroid-only distance
    int d = 64;
    int nlist = 8;
    int nb = 1000;
    int nq = 10;
    int k = 5;

    std::mt19937 rng(42);
    std::uniform_real_distribution<float> distrib(0.0f, 1.0f);
    std::vector<float> xb(nb * d), xq(nq * d);
    for (int i = 0; i < nb * d; i++) {
        xb[i] = distrib(rng);
    }
    for (int i = 0; i < nq * d; i++) {
        xq[i] = distrib(rng);
    }

    faiss::IndexFlatL2 quantizer(d);
    faiss::IndexIVFScalarQuantizer index(
            &quantizer,
            d,
            nlist,
            faiss::ScalarQuantizer::QT_0bit,
            faiss::METRIC_L2,
            false);
    EXPECT_EQ(index.code_size, 0);
    EXPECT_FALSE(index.by_residual);

    index.train(nb, xb.data());
    index.add(nb, xb.data());
    EXPECT_EQ(index.ntotal, nb);

    index.nprobe = nlist;
    std::vector<float> distances(nq * k);
    std::vector<faiss::idx_t> labels(nq * k);
    index.search(nq, xq.data(), k, distances.data(), labels.data());

    // Verify we got results
    for (int q = 0; q < nq; q++) {
        EXPECT_GE(labels[q * k], 0);
    }

    // Compare with direct quantizer search - distances should match
    std::vector<float> coarse_dis(nq * nlist);
    std::vector<faiss::idx_t> coarse_ids(nq * nlist);
    quantizer.search(
            nq, xq.data(), nlist, coarse_dis.data(), coarse_ids.data());

    for (int q = 0; q < nq; q++) {
        float ivf_dis = distances[q * k];
        bool found = false;
        for (int j = 0; j < nlist; j++) {
            if (std::abs(ivf_dis - coarse_dis[q * nlist + j]) < 1e-5) {
                found = true;
                break;
            }
        }
        EXPECT_TRUE(found) << "IVF distance " << ivf_dis
                           << " not found in coarse distances for query " << q;
    }
}

TEST(TestSQ0bit, IndexFactory) {
    int d = 32;
    std::unique_ptr<faiss::Index> index(faiss::index_factory(d, "IVF8,SQ0"));
    EXPECT_NE(index, nullptr);
    auto* ivfsq = dynamic_cast<faiss::IndexIVFScalarQuantizer*>(index.get());
    EXPECT_NE(ivfsq, nullptr);
    EXPECT_EQ(ivfsq->sq.qtype, faiss::ScalarQuantizer::QT_0bit);
    EXPECT_EQ(ivfsq->code_size, 0);
}

TEST(TestSQ0bit, InnerProduct) {
    int d = 64;
    int nlist = 4;
    int nb = 500;
    int nq = 5;
    int k = 3;

    std::mt19937 rng2(43);
    std::uniform_real_distribution<float> distrib2(0.0f, 1.0f);
    std::vector<float> xb(nb * d), xq(nq * d);
    for (int i = 0; i < nb * d; i++) {
        xb[i] = distrib2(rng2);
    }
    for (int i = 0; i < nq * d; i++) {
        xq[i] = distrib2(rng2);
    }

    faiss::IndexFlatIP quantizer(d);
    faiss::IndexIVFScalarQuantizer index(
            &quantizer,
            d,
            nlist,
            faiss::ScalarQuantizer::QT_0bit,
            faiss::METRIC_INNER_PRODUCT,
            false);

    index.train(nb, xb.data());
    index.add(nb, xb.data());

    index.nprobe = nlist;
    std::vector<float> distances(nq * k);
    std::vector<faiss::idx_t> labels(nq * k);
    index.search(nq, xq.data(), k, distances.data(), labels.data());

    for (int q = 0; q < nq; q++) {
        EXPECT_GE(labels[q * k], 0);
    }
}

TEST(ScalarQuantizer, TQMSEEncodeDecode) {
    check_tqmse_roundtrip(32, faiss::ScalarQuantizer::QT_1bit_tqmse);
    check_tqmse_roundtrip(32, faiss::ScalarQuantizer::QT_2bit_tqmse);
    check_tqmse_roundtrip(32, faiss::ScalarQuantizer::QT_3bit_tqmse);
    check_tqmse_roundtrip(32, faiss::ScalarQuantizer::QT_4bit_tqmse);
    check_tqmse_roundtrip(32, faiss::ScalarQuantizer::QT_8bit_tqmse);
}

TEST(ScalarQuantizer, EDENEncodeDecode) {
    check_eden_roundtrip(32, 1, faiss::ScalarQuantizer::QT_1bit_eden);
    check_eden_roundtrip(32, 2, faiss::ScalarQuantizer::QT_2bit_eden);
    check_eden_roundtrip(32, 3, faiss::ScalarQuantizer::QT_3bit_eden);
    check_eden_roundtrip(32, 4, faiss::ScalarQuantizer::QT_4bit_eden);
    check_eden_roundtrip(32, 5, faiss::ScalarQuantizer::QT_5bit_eden);
    check_eden_roundtrip(32, 6, faiss::ScalarQuantizer::QT_6bit_eden);
    check_eden_roundtrip(32, 7, faiss::ScalarQuantizer::QT_7bit_eden);
    check_eden_roundtrip(32, 8, faiss::ScalarQuantizer::QT_8bit_eden);
}

TEST(ScalarQuantizer, TQMSEAccuracyOrdering) {
    const size_t d = 32;
    const size_t n = 256;
    std::vector<float> x = make_normalized_vectors(n, d);
    const std::array<faiss::ScalarQuantizer::QuantizerType, 5> qtypes = {
            faiss::ScalarQuantizer::QT_1bit_tqmse,
            faiss::ScalarQuantizer::QT_2bit_tqmse,
            faiss::ScalarQuantizer::QT_3bit_tqmse,
            faiss::ScalarQuantizer::QT_4bit_tqmse,
            faiss::ScalarQuantizer::QT_8bit_tqmse};

    std::vector<float> mses;
    mses.reserve(qtypes.size());

    for (auto qtype : qtypes) {
        faiss::ScalarQuantizer sq(d, qtype);
        sq.train(n, x.data());

        std::vector<uint8_t> codes(sq.code_size * n, 0);
        sq.compute_codes(x.data(), codes.data(), n);

        std::vector<float> decoded(n * d);
        sq.decode(codes.data(), decoded.data(), n);
        mses.push_back(mean_squared_error(x, decoded));
    }

    for (size_t i = 0; i + 1 < mses.size(); ++i) {
        EXPECT_GE(mses[i], mses[i + 1]);
    }
}

TEST(ScalarQuantizer, TQMSENonSimdDims) {
    check_tqmse_roundtrip(7, faiss::ScalarQuantizer::QT_1bit_tqmse);
    check_tqmse_roundtrip(9, faiss::ScalarQuantizer::QT_2bit_tqmse);
    check_tqmse_roundtrip(11, faiss::ScalarQuantizer::QT_3bit_tqmse);
    check_tqmse_roundtrip(13, faiss::ScalarQuantizer::QT_4bit_tqmse);
    check_tqmse_roundtrip(33, faiss::ScalarQuantizer::QT_8bit_tqmse);
}

TEST(ScalarQuantizer, TQMSESimdDispatchSelection) {
    const std::vector<faiss::SIMDLevel> levels =
            available_lloyd_max_simd_levels();
    if (levels.empty()) {
        GTEST_SKIP() << "No SIMD level available for TurboQuant dispatch tests";
    }

    for (faiss::SIMDLevel level : levels) {
        SCOPED_TRACE(faiss::to_string(level));
        expect_lloyd_max_simd_dispatch_for_compatible_dim<1>(
                level, faiss::ScalarQuantizer::QT_1bit_tqmse);
        expect_lloyd_max_simd_dispatch_for_compatible_dim<2>(
                level, faiss::ScalarQuantizer::QT_2bit_tqmse);
        expect_lloyd_max_simd_dispatch_for_compatible_dim<3>(
                level, faiss::ScalarQuantizer::QT_3bit_tqmse);
        expect_lloyd_max_simd_dispatch_for_compatible_dim<4>(
                level, faiss::ScalarQuantizer::QT_4bit_tqmse);
        expect_lloyd_max_simd_dispatch_for_compatible_dim<8>(
                level, faiss::ScalarQuantizer::QT_8bit_tqmse);

        expect_lloyd_max_simd_dispatch_fallback_for_incompatible_dim<1>(
                level, 7, faiss::ScalarQuantizer::QT_1bit_tqmse);
        expect_lloyd_max_simd_dispatch_fallback_for_incompatible_dim<2>(
                level, 9, faiss::ScalarQuantizer::QT_2bit_tqmse);
        expect_lloyd_max_simd_dispatch_fallback_for_incompatible_dim<3>(
                level, 11, faiss::ScalarQuantizer::QT_3bit_tqmse);
        expect_lloyd_max_simd_dispatch_fallback_for_incompatible_dim<4>(
                level, 13, faiss::ScalarQuantizer::QT_4bit_tqmse);
        expect_lloyd_max_simd_dispatch_fallback_for_incompatible_dim<8>(
                level, 33, faiss::ScalarQuantizer::QT_8bit_tqmse);
    }
}

TEST(ScalarQuantizer, EDENSimdDispatchSelection) {
    const std::vector<faiss::SIMDLevel> levels =
            available_lloyd_max_simd_levels();
    if (levels.empty()) {
        GTEST_SKIP() << "No SIMD level available for EDEN dispatch tests";
    }

    for (faiss::SIMDLevel level : levels) {
        SCOPED_TRACE(faiss::to_string(level));
        expect_lloyd_max_simd_dispatch_for_compatible_dim<1>(
                level, faiss::ScalarQuantizer::QT_1bit_eden);
        expect_lloyd_max_simd_dispatch_for_compatible_dim<2>(
                level, faiss::ScalarQuantizer::QT_2bit_eden);
        expect_lloyd_max_simd_dispatch_for_compatible_dim<3>(
                level, faiss::ScalarQuantizer::QT_3bit_eden);
        expect_lloyd_max_simd_dispatch_for_compatible_dim<4>(
                level, faiss::ScalarQuantizer::QT_4bit_eden);
        expect_lloyd_max_simd_dispatch_for_compatible_dim<8>(
                level, faiss::ScalarQuantizer::QT_8bit_eden);

        expect_lloyd_max_simd_dispatch_fallback_for_incompatible_dim<1>(
                level, 7, faiss::ScalarQuantizer::QT_1bit_eden);
        expect_lloyd_max_simd_dispatch_fallback_for_incompatible_dim<2>(
                level, 9, faiss::ScalarQuantizer::QT_2bit_eden);
        expect_lloyd_max_simd_dispatch_fallback_for_incompatible_dim<3>(
                level, 11, faiss::ScalarQuantizer::QT_3bit_eden);
        expect_lloyd_max_simd_dispatch_fallback_for_incompatible_dim<4>(
                level, 13, faiss::ScalarQuantizer::QT_4bit_eden);
        expect_lloyd_max_simd_dispatch_fallback_for_incompatible_dim<8>(
                level, 33, faiss::ScalarQuantizer::QT_8bit_eden);
    }
}

TEST(ScalarQuantizer, TQMSESimdDistancePathParity) {
    const std::vector<faiss::SIMDLevel> levels =
            available_sq_parity_simd_levels();
    if (levels.empty()) {
        GTEST_SKIP() << "No SIMD level available for TurboQuant parity tests";
    }

    for (faiss::SIMDLevel level : levels) {
        SCOPED_TRACE(faiss::to_string(level));
        check_lloyd_max_distance_path_parity<1>(
                level, faiss::ScalarQuantizer::QT_1bit_tqmse);
        check_lloyd_max_distance_path_parity<2>(
                level, faiss::ScalarQuantizer::QT_2bit_tqmse);
        check_lloyd_max_distance_path_parity<3>(
                level, faiss::ScalarQuantizer::QT_3bit_tqmse);
        check_lloyd_max_distance_path_parity<4>(
                level, faiss::ScalarQuantizer::QT_4bit_tqmse);
        check_lloyd_max_distance_path_parity<8>(
                level, faiss::ScalarQuantizer::QT_8bit_tqmse);
    }
}

TEST(ScalarQuantizer, EDENSimdDistancePathParity) {
    const std::vector<faiss::SIMDLevel> levels =
            available_sq_parity_simd_levels();
    if (levels.empty()) {
        GTEST_SKIP() << "No SIMD level available for EDEN parity tests";
    }

    for (faiss::SIMDLevel level : levels) {
        SCOPED_TRACE(faiss::to_string(level));
        check_lloyd_max_distance_path_parity<1>(
                level, faiss::ScalarQuantizer::QT_1bit_eden);
        check_lloyd_max_distance_path_parity<2>(
                level, faiss::ScalarQuantizer::QT_2bit_eden);
        check_lloyd_max_distance_path_parity<3>(
                level, faiss::ScalarQuantizer::QT_3bit_eden);
        check_lloyd_max_distance_path_parity<4>(
                level, faiss::ScalarQuantizer::QT_4bit_eden);
        check_lloyd_max_distance_path_parity<8>(
                level, faiss::ScalarQuantizer::QT_8bit_eden);
    }
}

// RVV-versus-scalar parity for all quantizers with RVV kernels, over both
// metrics and dimensions around VLMAX boundaries. With e32m8 vectors,
// VLMAX = VLEN / 4: 32 lanes on QEMU's default VLEN=128 and up to 512
// lanes on VLEN=2048 hardware, so the dimensions below straddle those
// chunk boundaries as well as tiny tails.
// QT_fp16 is covered by the dedicated RVVFP16DistancePathParity test below
// so that the Zvfhmin-dependent case is visible individually in test logs.
TEST(ScalarQuantizer, RVVDistancePathParity) {
    if (!faiss::SIMDConfig::is_simd_level_available(
                faiss::SIMDLevel::RISCV_RVV)) {
        GTEST_SKIP() << "RISCV_RVV not available on this machine";
    }

    const std::vector<faiss::ScalarQuantizer::QuantizerType> qtypes = {
            faiss::ScalarQuantizer::QT_8bit,
            faiss::ScalarQuantizer::QT_4bit,
            faiss::ScalarQuantizer::QT_6bit,
            faiss::ScalarQuantizer::QT_8bit_uniform,
            faiss::ScalarQuantizer::QT_4bit_uniform,
            faiss::ScalarQuantizer::QT_bf16,
            faiss::ScalarQuantizer::QT_8bit_direct,
            faiss::ScalarQuantizer::QT_8bit_direct_signed,
            faiss::ScalarQuantizer::QT_1bit_eden,
            faiss::ScalarQuantizer::QT_2bit_eden,
            faiss::ScalarQuantizer::QT_3bit_eden,
            faiss::ScalarQuantizer::QT_4bit_eden,
            faiss::ScalarQuantizer::QT_8bit_eden,
            // The QT_*_tqmse types currently alias QuantizerLloydMax, so
            // they share the EDEN RVV kernels; listing them makes that
            // dispatch coverage explicit for every metric and distance API.
            faiss::ScalarQuantizer::QT_1bit_tqmse,
            faiss::ScalarQuantizer::QT_2bit_tqmse,
            faiss::ScalarQuantizer::QT_3bit_tqmse,
            faiss::ScalarQuantizer::QT_4bit_tqmse,
            faiss::ScalarQuantizer::QT_8bit_tqmse,
    };
    const std::vector<size_t> dims = {
            1,
            7,
            31,
            32,
            33,
            63,
            64,
            65,
            127,
            128,
            129,
            255,
            256,
            257,
            511,
            512,
            513};
    const std::vector<faiss::MetricType> metrics = {
            faiss::METRIC_L2, faiss::METRIC_INNER_PRODUCT};

    for (auto qtype : qtypes) {
        for (size_t d : dims) {
            for (auto metric : metrics) {
                SCOPED_TRACE(
                        testing::Message()
                        << "qtype=" << static_cast<int>(qtype) << " d=" << d
                        << " metric=" << metric);
                check_sq_distance_path_parity(
                        faiss::SIMDLevel::RISCV_RVV, qtype, metric, d);
            }
        }
    }
}

// QT_fp16 RVV-versus-scalar parity, split out of RVVDistancePathParity so
// the Zvfhmin-dependent case shows up individually in test logs. The RVV
// FP16 kernel emits Zvfhmin instructions, so this test only runs where the
// CPU/emulator enables the extension (CI uses -cpu rv64,v=true,
// x-zvfhmin=true); without it the process would die with SIGILL, which is
// what makes this test a check that the FP16 case genuinely executes.
TEST(ScalarQuantizer, RVVFP16DistancePathParity) {
    if (!faiss::SIMDConfig::is_simd_level_available(
                faiss::SIMDLevel::RISCV_RVV)) {
        GTEST_SKIP() << "RISCV_RVV not available on this machine";
    }

    const std::vector<size_t> dims = {
            1,
            7,
            31,
            32,
            33,
            63,
            64,
            65,
            127,
            128,
            129,
            255,
            256,
            257,
            511,
            512,
            513};
    const std::vector<faiss::MetricType> metrics = {
            faiss::METRIC_L2, faiss::METRIC_INNER_PRODUCT};

    for (size_t d : dims) {
        for (auto metric : metrics) {
            SCOPED_TRACE(
                    testing::Message() << "d=" << d << " metric=" << metric);
            check_sq_distance_path_parity(
                    faiss::SIMDLevel::RISCV_RVV,
                    faiss::ScalarQuantizer::QT_fp16,
                    metric,
                    d);
        }
    }
}

// Zero-dimensional regression for the d == 0 early returns in the RVV
// DCTemplate distance loops: query_to_code, query_to_codes_batch_4 and
// symmetric_dis must return exactly 0 without touching the query pointer or
// the codes, matching the scalar (NONE) path. The parity dimension lists
// above start at 1, so this case is covered here directly.
TEST(ScalarQuantizer, RVVZeroDimDistancePathParity) {
    if (!faiss::SIMDConfig::is_simd_level_available(
                faiss::SIMDLevel::RISCV_RVV)) {
        GTEST_SKIP() << "RISCV_RVV not available on this machine";
    }
    ScopedSIMDLevel scoped(faiss::SIMDLevel::RISCV_RVV);

    // trained layouts each quantizer expects at d == 0: non-uniform
    // templates take 2 * d (empty) values, uniform templates take {vmin,
    // vdiff}, LloydMax takes 2^k - 1 centroids/boundaries, and the raw
    // codecs (fp16/bf16/direct) ignore trained entirely.
    const std::vector<std::pair<
            faiss::ScalarQuantizer::QuantizerType,
            std::vector<float>>>
            cases = {
                    {faiss::ScalarQuantizer::QT_8bit, {}},
                    {faiss::ScalarQuantizer::QT_4bit, {}},
                    {faiss::ScalarQuantizer::QT_6bit, {}},
                    {faiss::ScalarQuantizer::QT_8bit_uniform, {0.0f, 1.0f}},
                    {faiss::ScalarQuantizer::QT_4bit_uniform, {0.0f, 1.0f}},
                    {faiss::ScalarQuantizer::QT_fp16, {}},
                    {faiss::ScalarQuantizer::QT_bf16, {}},
                    {faiss::ScalarQuantizer::QT_8bit_direct, {}},
                    {faiss::ScalarQuantizer::QT_8bit_direct_signed, {}},
                    {faiss::ScalarQuantizer::QT_1bit_eden,
                     std::vector<float>(3, 0.0f)},
                    {faiss::ScalarQuantizer::QT_2bit_eden,
                     std::vector<float>(7, 0.0f)},
                    {faiss::ScalarQuantizer::QT_3bit_eden,
                     std::vector<float>(15, 0.0f)},
                    {faiss::ScalarQuantizer::QT_4bit_eden,
                     std::vector<float>(31, 0.0f)},
                    {faiss::ScalarQuantizer::QT_8bit_eden,
                     std::vector<float>(511, 0.0f)},
            };
    const std::vector<faiss::MetricType> metrics = {
            faiss::METRIC_L2, faiss::METRIC_INNER_PRODUCT};

    uint8_t dummy_code[8] = {};

    for (const auto& [qtype, trained] : cases) {
        for (auto metric : metrics) {
            SCOPED_TRACE(
                    testing::Message() << "qtype=" << static_cast<int>(qtype)
                                       << " metric=" << metric);

            faiss::ScalarQuantizer sq(0, qtype);
            sq.trained = trained;

            std::unique_ptr<faiss::ScalarQuantizer::SQDistanceComputer>
                    scalar_dc(
                            faiss::scalar_quantizer::
                                    sq_select_distance_computer<
                                            faiss::SIMDLevel::NONE>(
                                            metric, qtype, 0, sq.trained));
            std::unique_ptr<faiss::ScalarQuantizer::SQDistanceComputer> simd_dc(
                    sq.get_distance_computer(metric));
            ASSERT_NE(scalar_dc, nullptr);
            ASSERT_NE(simd_dc, nullptr);

            scalar_dc->set_query(nullptr);
            simd_dc->set_query(nullptr);

            EXPECT_FLOAT_EQ(scalar_dc->query_to_code(dummy_code), 0.0f);
            EXPECT_FLOAT_EQ(simd_dc->query_to_code(dummy_code), 0.0f);

            float scalar_dis[4], simd_dis[4];
            scalar_dc->query_to_codes_batch_4(
                    dummy_code,
                    dummy_code,
                    dummy_code,
                    dummy_code,
                    scalar_dis[0],
                    scalar_dis[1],
                    scalar_dis[2],
                    scalar_dis[3]);
            simd_dc->query_to_codes_batch_4(
                    dummy_code,
                    dummy_code,
                    dummy_code,
                    dummy_code,
                    simd_dis[0],
                    simd_dis[1],
                    simd_dis[2],
                    simd_dis[3]);
            for (int k = 0; k < 4; k++) {
                EXPECT_FLOAT_EQ(scalar_dis[k], 0.0f);
                EXPECT_FLOAT_EQ(simd_dis[k], 0.0f);
            }

            scalar_dc->codes = dummy_code;
            scalar_dc->code_size = 0;
            simd_dc->codes = dummy_code;
            simd_dc->code_size = 0;
            EXPECT_FLOAT_EQ(scalar_dc->symmetric_dis(0, 1), 0.0f);
            EXPECT_FLOAT_EQ(simd_dc->symmetric_dis(0, 1), 0.0f);
        }
    }
}