/*
 * 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 <faiss/MetricType.h>
#include <faiss/impl/platform_macros.h>
#include <faiss/utils/AlignedTable.h>
#include <faiss/utils/rabitq_simd.h>
#include <faiss/utils/simd_levels.h>
#include <algorithm>
#include <cstddef>
#include <cstdint>
#include <cstring>
#include <vector>

namespace faiss {
namespace rabitq_utils {

/** Base factors computed per database vector for RaBitQ distance computation.
 * Used by both 1-bit and multi-bit RaBitQ variants.
 * These can be stored either embedded in codes (IndexRaBitQ) or separately
 * (IndexRaBitQFastScan).
 *
 * For 1-bit mode only - contains the minimal factors needed for distance
 * estimation using just sign bits.
 */
FAISS_PACK_STRUCTS_BEGIN
struct FAISS_PACKED SignBitFactors {
    // ||or - c||^2 - ((metric==IP) ? ||or||^2 : 0)
    float or_minus_c_l2sqr = 0;
    float dp_multiplier = 0;
};

/** Extended factors for multi-bit RaBitQ (nb_bits > 1).
 * Includes error bound for lower bound computation in two-stage search.
 * Inherits base factors to maintain layout compatibility.
 *
 * Used in multi-bit mode - the error bound enables quick filtering of
 * unlikely candidates in the first stage of two-stage search.
 */
struct FAISS_PACKED SignBitFactorsWithError : SignBitFactors {
    // Error bound for lower bound computation in two-stage search
    // Used in formula: lower_bound = est_distance - f_error * g_error
    // Only allocated when nb_bits > 1
    float f_error = 0;
};

/** Additional factors for multi-bit RaBitQ (nb_bits > 1).
 * Used to store normalization and scaling factors for the refinement bits
 * that encode additional precision beyond the sign bit.
 */
struct FAISS_PACKED ExtraBitsFactors {
    // Additive correction factor for refinement bit reconstruction
    float f_add_ex = 0;
    // Scaling/rescaling factor for refinement bit reconstruction
    float f_rescale_ex = 0;
};
FAISS_PACK_STRUCTS_END

/** Query-specific factors computed during search for RaBitQ distance
 * computation. Used by both IndexRaBitQ and IndexRaBitQFastScan
 * implementations.
 */
struct QueryFactorsData {
    float c1 = 0;
    float c2 = 0;
    float c34 = 0;

    float qr_to_c_L2sqr = 0;
    float qr_norm_L2sqr = 0;
    float q_dot_c = 0; // <query, centroid> for IP metric; 0 for L2

    float int_dot_scale = 1;

    float g_error = 0;
    std::vector<float> rotated_q;
};

/** Ideal quantizer radii for quantizers of 1..8 bits, optimized to minimize
 * L2 reconstruction error. Shared between all RaBitQ implementations.
 */
FAISS_API extern const float Z_MAX_BY_QB[8];

/** Compute factors for a single database vector using RaBitQ algorithm.
 * This function consolidates the mathematical logic that was duplicated
 * between IndexRaBitQ and IndexRaBitQFastScan.
 *
 * @param x             input vector (d dimensions)
 * @param d             dimensionality
 * @param centroid      database centroid (nullptr if not used)
 * @param metric_type   distance metric (L2 or Inner Product)
 * @param compute_error whether to compute f_error (false for 1-bit mode)
 * @return              computed factors for distance computation
 */
SignBitFactorsWithError compute_vector_factors(
        const float* x,
        size_t d,
        const float* centroid,
        MetricType metric_type,
        bool compute_error = true);

/** Compute intermediate values needed for vector factor computation.
 * Separated out to allow different bit packing strategies while sharing
 * the core mathematical computation.
 *
 * @param x             input vector (d dimensions)
 * @param d             dimensionality
 * @param centroid      database centroid (nullptr if not used)
 * @param norm_L2sqr    output: ||or - c||^2
 * @param or_L2sqr      output: ||or||^2
 * @param dp_oO         output: sum of |or_i - c_i| (absolute deviations)
 */
void compute_vector_intermediate_values(
        const float* x,
        size_t d,
        const float* centroid,
        float& norm_L2sqr,
        float& or_L2sqr,
        float& dp_oO);

/** Compute final factors from intermediate values.
 * @param norm_L2sqr    ||or - c||^2
 * @param or_L2sqr      ||or||^2
 * @param dp_oO         sum of |or_i - c_i|
 * @param d             dimensionality
 * @param metric_type   distance metric
 * @param compute_error whether to compute f_error (false for 1-bit mode)
 * @return              computed factors
 */
SignBitFactorsWithError compute_factors_from_intermediates(
        float norm_L2sqr,
        float or_L2sqr,
        float dp_oO,
        size_t d,
        MetricType metric_type,
        bool compute_error = true);

/** Compute query factors for RaBitQ distance computation.
 * This consolidates the query processing logic shared between implementations.
 *
 * @param query         query vector (d dimensions)
 * @param d             dimensionality
 * @param centroid      database centroid (nullptr if not used)
 * @param qb            number of quantization bits (1-8)
 * @param centered      whether to use centered quantization
 * @param metric_type   distance metric
 * @param rotated_q     output: query - centroid
 * @param rotated_qq    output: quantized query values
 * @return              computed query factors
 */
QueryFactorsData compute_query_factors(
        const float* query,
        size_t d,
        const float* centroid,
        uint8_t qb,
        bool centered,
        MetricType metric_type,
        std::vector<float>& rotated_q,
        std::vector<uint8_t>& rotated_qq);

/** Extract bit value from RaBitQ code in standard format.
 * Used by IndexRaBitQ which stores bits sequentially.
 *
 * @param code          RaBitQ code data
 * @param bit_index     which bit to extract (0 to d-1)
 * @return              bit value (true/false)
 */
bool extract_bit_standard(const uint8_t* code, size_t bit_index);

/** Extract bit value from FastScan code format.
 * Used by IndexRaBitQFastScan which packs bits into 4-bit sub-quantizers.
 *
 * @param code          FastScan code data
 * @param bit_index     which bit to extract (0 to d-1)
 * @return              bit value (true/false)
 */
bool extract_bit_fastscan(const uint8_t* code, size_t bit_index);

/** Set bit value in standard RaBitQ code format.
 * @param code          RaBitQ code data to modify
 * @param bit_index     which bit to set (0 to d-1)
 */
void set_bit_standard(uint8_t* code, size_t bit_index);

/** Set bit value in FastScan code format.
 * @param code          FastScan code data to modify
 * @param bit_index     which bit to set (0 to d-1)
 */
void set_bit_fastscan(uint8_t* code, size_t bit_index);

/** Compute adjusted 1-bit distance from normalized LUT distance.
 * This is the core distance formula shared by all RaBitQ handlers.
 *
 * @param normalized_distance  Distance from SIMD LUT lookup (after
 * normalization)
 * @param db_factors          Database vector factors (SignBitFactors or
 * SignBitFactorsWithError)
 * @param query_factors       Query factors computed during search
 * @param centered            Whether centered quantization is used
 * @param qb                  Number of quantization bits
 * @param d                   Dimensionality
 * @return                    Adjusted distance value
 */
inline float compute_1bit_adjusted_distance(
        float normalized_distance,
        const SignBitFactors& db_factors,
        const QueryFactorsData& query_factors,
        bool centered,
        size_t qb,
        size_t d) {
    float adjusted_distance;

    if (centered) {
        // For centered mode: normalized_distance contains the raw XOR
        // contribution. Apply the signed odd integer quantization formula:
        // int_dot = ((1 << qb) - 1) * d - 2 * xor_dot_product
        int64_t int_dot = ((1 << qb) - 1) * d;
        int_dot -= 2 * static_cast<int64_t>(normalized_distance);

        adjusted_distance = query_factors.qr_to_c_L2sqr +
                db_factors.or_minus_c_l2sqr -
                2 * db_factors.dp_multiplier * int_dot *
                        query_factors.int_dot_scale;
    } else {
        // For non-centered quantization: use traditional formula
        float final_dot = normalized_distance - query_factors.c34;
        adjusted_distance = db_factors.or_minus_c_l2sqr +
                query_factors.qr_to_c_L2sqr -
                2 * db_factors.dp_multiplier * final_dot;
    }

    // Apply inner product correction if needed
    if (query_factors.qr_norm_L2sqr != 0.0f) {
        adjusted_distance =
                -0.5f * (adjusted_distance - query_factors.qr_norm_L2sqr);
    } else {
        adjusted_distance = std::max(0.0f, adjusted_distance);
    }

    return adjusted_distance;
}

/** Determine whether a candidate should be refined in two-stage search.
 * Consolidates the filtering logic for both L2 and IP metrics.
 *
 * For L2 (min-heap): uses lower_bound = est_distance - error_adjustment
 *   - Skip if lower_bound >= threshold (can't beat current worst)
 * For IP (max-heap): uses upper_bound = est_distance + error_adjustment
 *   - Skip if upper_bound <= threshold (can't beat current best)
 *
 * @param est_distance     Estimated 1-bit distance
 * @param f_error          Database vector error factor
 * @param g_error          Query vector error factor
 * @param threshold        Current heap threshold (worst result in heap)
 * @param is_similarity    True for IP metric (max-heap), false for L2
 * (min-heap)
 * @return                 True if candidate should be refined with full
 * multi-bit distance
 */
inline bool should_refine_candidate(
        float est_distance,
        float f_error,
        float g_error,
        float threshold,
        bool is_similarity) {
    float error_adjustment = f_error * g_error;
    if (is_similarity) {
        // IP (max-heap): use upper bound for filtering
        float upper_bound = est_distance + error_adjustment;
        return upper_bound > threshold;
    } else {
        // L2 (min-heap): use lower bound for filtering
        float lower_bound = std::max(0.0f, est_distance - error_adjustment);
        return lower_bound < threshold;
    }
}

/** Extract multi-bit code on-the-fly from packed ex-bit codes.
 * This inline function extracts a single code value without unpacking the
 * entire array, enabling efficient on-the-fly decoding during distance
 * computation.
 *
 * @param ex_code       packed ex-bit codes
 * @param index         which code to extract (0 to d-1)
 * @param ex_bits       number of bits per code (1-8)
 * @return              extracted code value in range [0, 2^ex_bits - 1]
 */
inline int extract_code_inline(
        const uint8_t* ex_code,
        size_t index,
        size_t ex_bits) {
    size_t bit_pos = index * ex_bits;
    int code_value = 0;

    // Extract ex_bits bits starting at bit_pos
    for (size_t bit = 0; bit < ex_bits; bit++) {
        size_t byte_idx = bit_pos / 8;
        size_t bit_idx = bit_pos % 8;

        if (ex_code[byte_idx] & (1 << bit_idx)) {
            code_value |= (1 << bit);
        }

        bit_pos++;
    }

    return code_value;
}

/** Compute full multi-bit distance from sign bits and ex-bit codes.
 * This is the core distance computation shared by RaBitQFastScan handlers.
 *
 * The multi-bit distance combines the sign bit (1-bit) with additional
 * magnitude bits (ex_bits) to compute a more accurate distance estimate.
 * Uses SIMD-optimized bit-plane decomposition (AVX2+BMI2) for ex_bits 1-7,
 * with scalar fallback for non-x86 or non-BMI2 platforms.
 *
 * @param sign_bits       unpacked sign bits (1-bit codes in standard format)
 * @param ex_code         packed ex-bit codes
 * @param ex_fac          ex-bit factors (f_add_ex, f_rescale_ex)
 * @param rotated_q       rotated query vector
 * @param qr_base         precomputed base term: ||q-c||^2 for L2, <q,c> for IP
 * @param d               dimensionality
 * @param ex_bits         number of extra bits (nb_bits - 1)
 * @param metric_type     distance metric (L2 or Inner Product)
 * @return                computed full multi-bit distance
 */
float compute_full_multibit_distance(
        const uint8_t* sign_bits,
        const uint8_t* ex_code,
        const ExtraBitsFactors& ex_fac,
        const float* rotated_q,
        float qr_base,
        size_t d,
        size_t ex_bits,
        MetricType metric_type);

// SIMDLevel-templatized version — avoids per-call dynamic dispatch.
// Inline so it can be used from templatized distance computers without
// needing explicit instantiations in per-SIMD TUs.
template <SIMDLevel SL>
inline float compute_full_multibit_distance(
        const uint8_t* sign_bits,
        const uint8_t* ex_code,
        const ExtraBitsFactors& ex_fac,
        const float* rotated_q,
        float qr_base,
        size_t d,
        size_t ex_bits,
        MetricType metric_type) {
    const float cb = -(static_cast<float>(1 << ex_bits) - 0.5f);

    float ex_ip = rabitq::multibit::compute_inner_product<SL>(
            sign_bits, ex_code, rotated_q, d, ex_bits, cb);

    float dist = qr_base + ex_fac.f_add_ex + ex_fac.f_rescale_ex * ex_ip;

    if (metric_type == MetricType::METRIC_L2) {
        dist = std::max(0.0f, dist);
    }

    return dist;
}

/** Compute pointer to a vector's auxiliary data within block layout. */
template <typename T>
inline T* get_block_aux_ptr(
        T* block_data,
        size_t vec_pos,
        size_t bbs,
        size_t packed_block_size,
        size_t full_block_size,
        size_t storage_size) {
    return block_data + (vec_pos / bbs) * full_block_size + packed_block_size +
            (vec_pos % bbs) * storage_size;
}

/// Extract sign bits from PQ4-interleaved block into flat byte packing.
/// Like CodePackerRaBitQ::unpack_1 but sign-bits-only and with the
/// vector's in-block address hoisted out of the per-SQ loop.
inline void unpack_sign_bits_from_packed(
        const uint8_t* block,
        size_t bbs,
        size_t nsq,
        size_t offset,
        size_t block_stride,
        uint8_t* sign_bits_out) {
    block += (offset / bbs) * block_stride;
    offset = offset % bbs;

    const bool nibble_high = offset > 15;
    const size_t vid = offset & 15;
    const size_t in_group_addr =
            (vid < 8) ? (vid << 1) : (((vid - 8) << 1) + 1);

    const size_t num_pairs = nsq / 2;
    for (size_t k = 0; k < num_pairs; k++) {
        const size_t base = k * bbs;
        const uint8_t raw_even = block[base + in_group_addr];
        const uint8_t raw_odd = block[base + in_group_addr + 16];

        const uint8_t nib0 = nibble_high ? (raw_even >> 4) : (raw_even & 0xF);
        const uint8_t nib1 = nibble_high ? (raw_odd >> 4) : (raw_odd & 0xF);
        sign_bits_out[k] = nib0 | (nib1 << 4);
    }

    if (nsq & 1) {
        const uint8_t raw = block[num_pairs * bbs + in_group_addr];
        sign_bits_out[num_pairs] = nibble_high ? (raw >> 4) : (raw & 0xF);
    }
}

/** Compute per-vector auxiliary storage size.
 *
 * @param nb_bits  number of quantization bits (1 = sign-bit only)
 * @param d        dimensionality
 * @return         storage size in bytes
 */
size_t compute_per_vector_storage_size(size_t nb_bits, size_t d);

/** [LEGACY FORMAT SUPPORT] Migrate block data from old I/O format to new
 * format.
 *
 * This function is used only when reading indexes saved with the legacy format
 * (fourcc "Irfs"/"Iwrf") to convert them to the new embedded auxiliary data
 * format. Not needed for indexes saved with the new format ("Irfn"/"Iwrn").
 *
 * Re-layouts blocks in-place and copies aux data from flat_storage.
 *
 * @param flat_storage       legacy per-vector aux data indexed by global ID
 * @param codes              block data (will be resized and re-laid out)
 * @param num_vectors        number of vectors in this segment
 * @param bbs                block batch size (vectors per block)
 * @param M2                 rounded sub-quantizer count
 * @param old_block_stride   old block size (packed codes only, or current)
 * @param new_block_stride   new block size (packed codes + aux region)
 * @param storage_size       per-vector aux storage size in bytes
 * @param id_map             maps local offset to global ID; null = sequential
 */
void populate_block_aux_from_flat_storage(
        const std::vector<uint8_t>& flat_storage,
        AlignedTable<uint8_t>& codes,
        size_t num_vectors,
        size_t bbs,
        size_t M2,
        size_t old_block_stride,
        size_t new_block_stride,
        size_t storage_size,
        const int64_t* id_map = nullptr);

} // namespace rabitq_utils
} // namespace faiss
