# Copyright (c) Facebook, Inc. and its affiliates.
#
# This source code is licensed under the MIT license found in the
# LICENSE file in the root directory of this source tree.

# @nolint

# not linting this file because it imports * from swigfaiss, which
# causes a ton of useless warnings.

import numpy as np
import logging
import sys
import inspect


def _preload_gpu_libs():
    """Pre-load CUDA / RAPIDS shared libs from pip wheels with RTLD_GLOBAL.

    These libs ship in nvidia-*-cuNN / libcuvs-cuNN wheels, off ld.so's search
    path, so we dlopen them before the SWIG extension loads libfaiss.so. Gated
    on the `faiss._gpu_build` marker (CMake writes it only for GPU builds); the
    `_cuvs_build` marker selects the CUDA 13 cuVS variant (else CUDA 12) and
    adds the cuVS stack. Missing wheels raise a fix-it.
    """
    try:
        from . import _gpu_build  # noqa: F401
    except ImportError:
        return  # faiss-cpu install: marker absent, nothing to preload

    import ctypes
    import os

    def _load(path):
        """dlopen one .so with global visibility, or raise a fix-it."""
        try:
            ctypes.CDLL(path, mode=ctypes.RTLD_GLOBAL)
        except OSError as e:
            raise RuntimeError(
                f"faiss-gpu: failed to load {os.path.basename(path)} from "
                f"{path} — corrupt or incomplete nvidia CUDA wheel?"
            ) from e

    # faiss-gpu-cuvs wheels carry the `_cuvs_build` marker and are built against
    # CUDA 13 (+ the cuVS stack); the plain faiss-gpu wheel is CUDA 12. The two
    # use different pip wheel layouts (handled separately below). Load order
    # matters: libcudart first (others resolve symbols against it), and load
    # RTLD_GLOBAL so the SWIG extension's later dlopen of libfaiss.so sees them.
    try:
        from . import _cuvs_build  # noqa: F401

        is_cuvs = True
    except ImportError:
        is_cuvs = False

    if is_cuvs:
        # CUDA 13 nvidia wheels (nvidia-cuda-runtime / -cublas / -curand /
        # -nvjitlink — all unsuffixed) ship every lib in one PEP 420 namespace
        # dir, nvidia/cu13/lib/, with no importable per-component module. This
        # differs from CUDA 12's per-component nvidia/<name>/lib/ layout.
        try:
            import nvidia.cu13 as _cu13

            _libdir = os.path.join(list(_cu13.__path__)[0], "lib")
        except ImportError as e:
            raise RuntimeError(
                "faiss-gpu-cuvs installed but the CUDA 13 runtime wheels are "
                "missing — pip install 'nvidia-cuda-runtime>=13.2,<14' "
                "'nvidia-cublas>=13.2,<14' 'nvidia-curand>=10.3.7,<11' "
                "'nvidia-nvjitlink>=13.2,<14'"
            ) from e
        # nvjitlink before cublas: cublas 13 links it.
        for _soname in (
            "libcudart.so.13",
            "libnvJitLink.so.13",
            "libcublas.so.13",
            "libcublasLt.so.13",
            "libcurand.so.10",
        ):
            _load(os.path.join(_libdir, _soname))
    else:
        # CUDA 12 per-component layout: each nvidia-*-cu12 wheel exposes an
        # importable module whose lib/ dir holds the .so.
        def _nvidia_lib_dir(import_name, pip_spec):
            """Return the lib/ dir of an nvidia-*-cu12 wheel, or raise a
            fix-it."""
            try:
                mod = __import__(
                    "nvidia." + import_name, fromlist=[import_name]
                )
            except ImportError as e:
                pip_name = pip_spec.split(">")[0].split("<")[0].split("=")[0]
                raise RuntimeError(
                    f"faiss-gpu installed but {pip_name} is missing — "
                    f"pip install '{pip_spec}'"
                ) from e
            # __path__[0] not __file__: PEP 420 namespace pkgs have
            # __file__ = None.
            return os.path.join(mod.__path__[0], "lib")

        _cudart = _nvidia_lib_dir(
            "cuda_runtime", "nvidia-cuda-runtime-cu12>=12.6,<13"
        )
        _cublas = _nvidia_lib_dir("cublas", "nvidia-cublas-cu12>=12.6,<13")
        _curand = _nvidia_lib_dir("curand", "nvidia-curand-cu12>=10.3.7,<11")
        _load(os.path.join(_cudart, "libcudart.so.12"))
        _load(os.path.join(_cublas, "libcublas.so.12"))
        _load(os.path.join(_cublas, "libcublasLt.so.12"))
        _load(os.path.join(_curand, "libcurand.so.10"))
        return  # faiss-gpu (non-cuVS) build: base CUDA preload is sufficient.

    # faiss-gpu-cuvs wheels also need the cuVS stack. Delegate to RAPIDS'
    # load_library() (loads each .so RTLD_GLOBAL + its CUDA deps); order
    # rmm -> raft -> cuvs makes every symbol global before the SWIG extension
    # loads.
    try:
        import libcuvs
        import libraft
        import librmm
    except ImportError as e:
        raise RuntimeError(
            "faiss-gpu-cuvs installed but the cuVS runtime wheels are "
            "missing — "
            "pip install 'libcuvs-cu13>=26.06,<27' "
            "--extra-index-url https://pypi.nvidia.com"
        ) from e

    for _mod in (librmm, libraft, libcuvs):
        _mod.load_library()


_preload_gpu_libs()
del _preload_gpu_libs

# We import * so that the symbol foo can be accessed as faiss.foo.
from .loader import *

# additional wrappers
from faiss import class_wrappers
from faiss.gpu_wrappers import *
from faiss.array_conversions import *
from faiss.extra_wrappers import (
    kmin,
    kmax,
    pairwise_distances,
    rand,
    randint,
    lrand,
    randn,
    rand_smooth_vectors,
    eval_intersection,
    normalize_L2,
    ResultHeap,
    knn,
    Kmeans,
    SuperKmeans,
    checksum,
    matrix_bucket_sort_inplace,
    bucket_sort,
    merge_knn_results,
    MapInt64ToInt64,
    knn_hamming,
    pack_bitstrings,
    unpack_bitstrings,
)


__version__ = "%d.%d.%d" % (
    FAISS_VERSION_MAJOR,
    FAISS_VERSION_MINOR,
    FAISS_VERSION_PATCH,
)

logger = logging.getLogger(__name__)

class_wrappers.handle_Clustering(Clustering)
class_wrappers.handle_SuperKMeans(SuperKMeans)
class_wrappers.handle_Clustering1D(Clustering1D)
class_wrappers.handle_MatrixStats(MatrixStats)
class_wrappers.handle_IOWriter(IOWriter)
class_wrappers.handle_IOReader(IOReader)
class_wrappers.handle_AutoTuneCriterion(AutoTuneCriterion)
class_wrappers.handle_ParameterSpace(ParameterSpace)
class_wrappers.handle_NSG(IndexNSG)
class_wrappers.handle_MapLong2Long(MapLong2Long)
class_wrappers.handle_IDSelectorSubset(IDSelectorBatch, class_owns=True)
class_wrappers.handle_IDSelectorSubset(IDSelectorArray, class_owns=False)
class_wrappers.handle_IDSelectorSubset(
    IDSelectorBitmap, class_owns=False, force_int64=False
)
class_wrappers.handle_CodeSet(CodeSet)

class_wrappers.handle_Tensor2D(Tensor2D)
class_wrappers.handle_Tensor2D(Int32Tensor2D)
class_wrappers.handle_Embedding(Embedding)
class_wrappers.handle_Linear(Linear)
class_wrappers.handle_QINCo(QINCo)
class_wrappers.handle_QINCoStep(QINCoStep)
shard_ivf_index_centroids = class_wrappers.handle_shard_ivf_index_centroids(
    shard_ivf_index_centroids
)


this_module = sys.modules[__name__]

# handle sub-classes
for symbol in dir(this_module):
    obj = getattr(this_module, symbol)
    # print symbol, isinstance(obj, (type, types.ClassType))
    if inspect.isclass(obj):
        the_class = obj
        if issubclass(the_class, Index):
            class_wrappers.handle_Index(the_class)

        if issubclass(the_class, IndexBinary):
            class_wrappers.handle_IndexBinary(the_class)

        if issubclass(the_class, VectorTransform):
            class_wrappers.handle_VectorTransform(the_class)

        if issubclass(the_class, Quantizer):
            class_wrappers.handle_Quantizer(the_class)

        if issubclass(the_class, IndexRowwiseMinMax) or issubclass(
            the_class, IndexRowwiseMinMaxFP16
        ):
            class_wrappers.handle_IndexRowwiseMinMax(the_class)

        if issubclass(the_class, SearchParameters):
            class_wrappers.handle_SearchParameters(the_class)

        if issubclass(the_class, CodePacker):
            class_wrappers.handle_CodePacker(the_class)

##############################################################################
# For some classes (IndexIVF, IDSelector), the object holds a reference to
# a C++ object (eg. the quantizer object of IndexIVF). We don't transfer the
# ownership to the C++ object (ie. set own_quantizer=true), but instead we add
# a reference in the Python class wrapper instead. This is done via an
# additional referenced_objects field.
#
# Since the semantics of ownership in the C++ classes are sometimes irregular,
# these references are added manually using the functions below.
##############################################################################


def add_ref_in_constructor(the_class, parameter_no):
    # adds a reference to parameter parameter_no in self
    # so that that parameter does not get deallocated before self
    original_init = the_class.__init__

    def replacement_init(self, *args):
        original_init(self, *args)
        self.referenced_objects = [args[parameter_no]]

    def replacement_init_multiple(self, *args):
        original_init(self, *args)
        pset = parameter_no[len(args)]
        self.referenced_objects = [args[no] for no in pset]

    if type(parameter_no) == dict:
        # a list of parameters to keep, depending on the number of arguments
        the_class.__init__ = replacement_init_multiple
    else:
        the_class.__init__ = replacement_init


def add_to_referenced_objects(self, ref):
    if not hasattr(self, "referenced_objects"):
        self.referenced_objects = [ref]
    else:
        self.referenced_objects.append(ref)


def add_ref_in_method(the_class, method_name, parameter_no):
    original_method = getattr(the_class, method_name)

    def replacement_method(self, *args):
        ref = args[parameter_no]
        add_to_referenced_objects(self, ref)
        return original_method(self, *args)

    setattr(the_class, method_name, replacement_method)


def add_ref_in_method_explicit_own(the_class, method_name):
    # for methods of format set_XXX(object, own)
    original_method = getattr(the_class, method_name)

    def replacement_method(self, ref, own=False):
        if not own:
            if not hasattr(self, "referenced_objects"):
                self.referenced_objects = [ref]
            else:
                self.referenced_objects.append(ref)
        else:
            # transfer ownership to C++ class
            ref.this.disown()
        return original_method(self, ref, own)

    setattr(the_class, method_name, replacement_method)


def add_ref_in_function(function_name, parameter_no):
    # assumes the function returns an object
    original_function = getattr(this_module, function_name)

    def replacement_function(*args):
        result = original_function(*args)
        ref = args[parameter_no]
        result.referenced_objects = [ref]
        return result

    setattr(this_module, function_name, replacement_function)


if "GPU" in get_compile_options():
    add_ref_in_constructor(GpuIndexIVFFlat, 1)
    add_ref_in_constructor(GpuIndexBinaryFlat, 1)
    add_ref_in_constructor(GpuIndexFlat, 1)
    add_ref_in_constructor(GpuIndexIVFPQ, 1)
    add_ref_in_constructor(GpuIndexIVFScalarQuantizer, 1)

add_ref_in_constructor(IndexIVFFlat, 0)
add_ref_in_constructor(IndexIVFFlatDedup, 0)
add_ref_in_constructor(IndexIVFFlatPanorama, 0)
add_ref_in_constructor(IndexPreTransform, {2: [0, 1], 1: [0]})
add_ref_in_method(IndexPreTransform, "prepend_transform", 0)
add_ref_in_constructor(IndexIVFPQ, 0)
add_ref_in_constructor(IndexIVFPQR, 0)
add_ref_in_constructor(IndexIVFPQFastScan, 0)
add_ref_in_constructor(IndexIVFResidualQuantizer, 0)
add_ref_in_constructor(IndexIVFLocalSearchQuantizer, 0)
add_ref_in_constructor(IndexIVFResidualQuantizerFastScan, 0)
add_ref_in_constructor(IndexIVFLocalSearchQuantizerFastScan, 0)
add_ref_in_constructor(IndexIVFSpectralHash, 0)
add_ref_in_method_explicit_own(IndexIVFSpectralHash, "replace_vt")

add_ref_in_constructor(Index2Layer, 0)
add_ref_in_constructor(Level1Quantizer, 0)
add_ref_in_constructor(IndexIVFScalarQuantizer, 0)
add_ref_in_constructor(IndexRowwiseMinMax, 0)
add_ref_in_constructor(IndexRowwiseMinMaxFP16, 0)
add_ref_in_constructor(IndexIDMap, 0)
add_ref_in_constructor(IndexIDMap2, 0)
add_ref_in_constructor(IndexHNSW, 0)
add_ref_in_method(IndexShards, "add_shard", 0)
add_ref_in_method(IndexBinaryShards, "add_shard", 0)
add_ref_in_constructor(IndexRefineFlat, {2: [0], 1: [0]})
add_ref_in_constructor(IndexRefinePanorama, {2: [0, 1]})
add_ref_in_constructor(IndexRefine, {2: [0, 1]})

add_ref_in_constructor(IndexBinaryIVF, 0)
add_ref_in_constructor(IndexBinaryFromFloat, 0)
add_ref_in_constructor(IndexBinaryIDMap, 0)
add_ref_in_constructor(IndexBinaryIDMap2, 0)

add_ref_in_method(IndexReplicas, "addIndex", 0)
add_ref_in_method(IndexBinaryReplicas, "addIndex", 0)

add_ref_in_constructor(BufferedIOWriter, 0)
add_ref_in_constructor(BufferedIOReader, 0)

add_ref_in_constructor(IDSelectorNot, 0)
add_ref_in_constructor(IDSelectorAnd, slice(2))
add_ref_in_constructor(IDSelectorOr, slice(2))
add_ref_in_constructor(IDSelectorXOr, slice(2))
add_ref_in_constructor(IDSelectorTranslated, slice(2))

add_ref_in_constructor(IDSelectorXOr, slice(2))
add_ref_in_constructor(IndexIVFIndependentQuantizer, slice(3))

add_ref_in_constructor(IndexIVFRaBitQ, 0)
add_ref_in_constructor(IndexIVFRaBitQFastScan, 0)
add_ref_in_constructor(IndexIVFEDEN, 0)

if "SVS" in get_compile_options():
    add_ref_in_constructor(IndexSVSVamana, 0)
    add_ref_in_constructor(IndexSVSVamanaLVQ, 0)
    add_ref_in_constructor(IndexSVSVamanaLeanVec, 0)
    add_ref_in_constructor(IndexSVSFlat, 0)

# seems really marginal...
# remove_ref_from_method(IndexReplicas, 'removeIndex', 0)


######################################################
# search_with_parameters interface
######################################################

search_with_parameters_c = search_with_parameters


def search_with_parameters(index, x, k, params=None, output_stats=False):
    x = np.ascontiguousarray(x, dtype="float32")
    n, d = x.shape
    assert d == index.d
    if not params:
        # if not provided use the ones set in the IVF object
        params = IVFSearchParameters()
        index_ivf = extract_index_ivf(index)
        params.nprobe = index_ivf.nprobe
        params.max_codes = index_ivf.max_codes
    nb_dis = np.empty(1, "uint64")
    ms_per_stage = np.empty(3, "float64")
    distances = np.empty((n, k), dtype=np.float32)
    labels = np.empty((n, k), dtype=np.int64)
    search_with_parameters_c(
        index,
        n,
        swig_ptr(x),
        k,
        swig_ptr(distances),
        swig_ptr(labels),
        params,
        swig_ptr(nb_dis),
        swig_ptr(ms_per_stage),
    )
    if not output_stats:
        return distances, labels
    else:
        stats = {
            "ndis": nb_dis[0],
            "pre_transform_ms": ms_per_stage[0],
            "coarse_quantizer_ms": ms_per_stage[1],
            "invlist_scan_ms": ms_per_stage[2],
        }
        return distances, labels, stats


range_search_with_parameters_c = range_search_with_parameters


def range_search_with_parameters(
    index, x, radius, params=None, output_stats=False
):
    x = np.ascontiguousarray(x, dtype="float32")
    n, d = x.shape
    assert d == index.d
    if not params:
        # if not provided use the ones set in the IVF object
        params = IVFSearchParameters()
        index_ivf = extract_index_ivf(index)
        params.nprobe = index_ivf.nprobe
        params.max_codes = index_ivf.max_codes
    nb_dis = np.empty(1, "uint64")
    ms_per_stage = np.empty(3, "float64")
    res = RangeSearchResult(n)
    range_search_with_parameters_c(
        index,
        n,
        swig_ptr(x),
        radius,
        res,
        params,
        swig_ptr(nb_dis),
        swig_ptr(ms_per_stage),
    )
    lims = rev_swig_ptr(res.lims, n + 1).copy()
    nd = int(lims[-1])
    Dout = rev_swig_ptr(res.distances, nd).copy()
    Iout = rev_swig_ptr(res.labels, nd).copy()
    if not output_stats:
        return lims, Dout, Iout
    else:
        stats = {
            "ndis": nb_dis[0],
            "pre_transform_ms": ms_per_stage[0],
            "coarse_quantizer_ms": ms_per_stage[1],
            "invlist_scan_ms": ms_per_stage[2],
        }
        return lims, Dout, Iout, stats


super_kmeans_assign_iteration_c = super_kmeans_assign_iteration


def super_kmeans_assign_iteration(
    X_tilde, Y_tilde, tau, assignments, d_prime, ad_coeff, cp,
):
    """Run one SuperKMeans iter-1+ pruned assignment pass on caller-managed state.

    All arrays must be C-contiguous. `X_tilde`, `Y_tilde`, `tau`, and `ad_coeff`
    must be float32; `assignments` must be int32. These mirror the C++ pointer
    contract, and passing another dtype (e.g. int64 assignments, numpy's default
    integer type) would otherwise reinterpret the buffer and corrupt results.

    Shapes: `X_tilde` is (n, d), `Y_tilde` is (k, d), `tau` and `assignments`
    have length n, and `ad_coeff` has length d + 1.

    Returns (total_pairs, pruned_at_gemm) for a d_prime controller.
    Mutates `tau` and `assignments` in place.
    """
    for name, arr, dtype in (
        ("X_tilde", X_tilde, "float32"),
        ("Y_tilde", Y_tilde, "float32"),
        ("tau", tau, "float32"),
        ("ad_coeff", ad_coeff, "float32"),
        ("assignments", assignments, "int32"),
    ):
        if arr.dtype != dtype:
            raise TypeError(
                f"super_kmeans_assign_iteration: {name} must be {dtype}, "
                f"got {arr.dtype}"
            )
        if not arr.flags["C_CONTIGUOUS"]:
            raise ValueError(
                f"super_kmeans_assign_iteration: {name} must be C-contiguous"
            )
    if X_tilde.ndim != 2:
        raise ValueError(
            f"super_kmeans_assign_iteration: X_tilde must be 2D (n, d), "
            f"got shape {X_tilde.shape}"
        )
    n, d = X_tilde.shape
    if Y_tilde.ndim != 2 or Y_tilde.shape[1] != d:
        raise ValueError(
            f"super_kmeans_assign_iteration: Y_tilde must have shape (k, {d}), "
            f"got {Y_tilde.shape}"
        )
    k = Y_tilde.shape[0]
    if tau.shape != (n,):
        raise ValueError(
            f"super_kmeans_assign_iteration: tau must have shape ({n},), "
            f"got {tau.shape}"
        )
    if assignments.shape != (n,):
        raise ValueError(
            f"super_kmeans_assign_iteration: assignments must have shape ({n},), "
            f"got {assignments.shape}"
        )
    if ad_coeff.shape != (d + 1,):
        raise ValueError(
            f"super_kmeans_assign_iteration: ad_coeff must have shape ({d + 1},), "
            f"got {ad_coeff.shape}"
        )
    total = np.zeros(1, dtype=np.int64)
    pruned = np.zeros(1, dtype=np.int64)
    super_kmeans_assign_iteration_c(
        swig_ptr(X_tilde), n, d,
        swig_ptr(Y_tilde), k,
        swig_ptr(tau), swig_ptr(assignments),
        d_prime, swig_ptr(ad_coeff), cp,
        swig_ptr(total), swig_ptr(pruned),
    )
    return int(total[0]), int(pruned[0])


# IndexProxy was renamed to IndexReplicas, remap the old name for any old code
# people may have
IndexProxy = IndexReplicas
ConcatenatedInvertedLists = HStackInvertedLists
IndexResidual = IndexResidualQuantizer

IVFSearchParameters = SearchParametersIVF

###########################################
# serialization of indexes to byte arrays
###########################################


def serialize_index(index, io_flags=0):
    """convert an index to a numpy uint8 array"""
    writer = VectorIOWriter()
    write_index(index, writer, io_flags)
    return vector_to_array(writer.data)


def deserialize_index(data, io_flags=0):
    reader = VectorIOReader()
    copy_array_to_vector(data, reader.data)
    return read_index(reader, io_flags)


def serialize_index_binary(index):
    """convert an index to a numpy uint8 array"""
    writer = VectorIOWriter()
    write_index_binary(index, writer)
    return vector_to_array(writer.data)


def deserialize_index_binary(data):
    reader = VectorIOReader()
    copy_array_to_vector(data, reader.data)
    return read_index_binary(reader)


class TimeoutGuard:
    def __init__(self, timeout_in_seconds: float):
        self.timeout = timeout_in_seconds

    def __enter__(self):
        TimeoutCallback.reset(self.timeout)

    def __exit__(self, exc_type, exc_value, traceback):
        PythonInterruptCallback.reset()


try:
    post_init_hook()
except NameError:
    pass