"""
Dense ANN module tests
"""

import os
import platform
import sys
import tempfile
import time
import unittest

from unittest.mock import patch

import ggml
import numpy as np

from sqlalchemy.dialects.postgresql import BIT
from sqlalchemy.ext.compiler import compiles

from txtai.ann import ANNFactory, ANN
from txtai.ann.dense.ggml import GGMLTensors
from txtai.serialize import SerializeFactory


class TestDense(unittest.TestCase):
    """
    Dense ANN tests.
    """

    def testAnnoy(self):
        """
        Test Annoy backend
        """

        self.runTests("annoy", None, False)

    def testAnnoyCustom(self):
        """
        Test Annoy backend with custom settings
        """

        # Test with custom settings
        self.runTests("annoy", {"annoy": {"ntrees": 2, "searchk": 1}}, False)

    def testCustomBackend(self):
        """
        Test resolving a custom backend
        """

        self.runTests("txtai.ann.Faiss")

    def testCustomBackendInvalid(self):
        """
        Test resolving an invalid backend
        """

        with self.assertRaises(ImportError):
            ANNFactory.create({"backend": "pprint.pprint"})

    def testCustomBackendNotFound(self):
        """
        Test resolving an unresolvable backend
        """

        with self.assertRaises(ImportError):
            ANNFactory.create({"backend": "notfound.ann"})

    def testFaiss(self):
        """
        Test Faiss backend
        """

        self.runTests("faiss")

    def testFaissBinary(self):
        """
        Test Faiss backend with a binary hash index
        """

        ann = ANNFactory.create({"backend": "faiss", "quantize": 1, "dimensions": 240 * 8, "faiss": {"components": "BHash32"}})

        # Generate and index dummy data
        data = np.random.rand(100, 240).astype(np.uint8)
        ann.index(data)

        # Generate query vector and test search
        query = np.random.rand(240).astype(np.uint8)
        self.assertGreater(ann.search(np.array([query]), 1)[0][0][1], 0)

    def testFaissCustom(self):
        """
        Test Faiss backend with custom settings
        """

        # Test with custom settings
        self.runTests("faiss", {"faiss": {"nprobe": 2, "components": "PCA16,IDMap,SQ8", "sample": 1.0}}, False)
        self.runTests("faiss", {"faiss": {"components": "IVF,SQ8"}}, False)

    @patch("platform.system")
    def testFaissMacOS(self, system):
        """
        Test Faiss backend with macOS
        """

        # Run test
        system.return_value = "Darwin"

        # pylint: disable=C0415, W0611
        # Force reload of class
        name = "txtai.ann.dense.faiss"
        module = sys.modules[name]
        del sys.modules[name]
        import txtai.ann.dense.faiss

        # Run tests
        self.runTests("faiss")

        # Restore original module
        sys.modules[name] = module

    @unittest.skipIf(os.name == "nt", "mmap not supported on Windows")
    def testFaissMmap(self):
        """
        Test Faiss backend with mmap enabled
        """

        # Test to with mmap enabled
        self.runTests("faiss", {"faiss": {"mmap": True}}, False)

    def testFaissNprobeClamp(self):
        """
        Test Faiss default nprobe clamp
        """

        ann, _ = self.faissmodel(5001, {"sample": 0.01})
        self.assertGreater(round(ann.cells(ann.count()) / 16), ann.backend.nlist)
        self.assertEqual(ann.nprobe(), ann.backend.nlist)

    def testFaissNprobeConfigured(self):
        """
        Test configured Faiss nprobe values
        """

        for nprobe in [1, 3]:
            with self.subTest(nprobe=nprobe):
                ann, _ = self.faissmodel(100, {"components": "IVF2,Flat", "nprobe": nprobe})
                self.assertEqual(ann.nprobe(), nprobe)

    def testFaissNprobeResults(self):
        """
        Test Faiss default nprobe clamp results
        """

        ann, rng = self.faissmodel(5001, {"sample": 0.01})
        queries = rng.random((10, 8), dtype=np.float32)
        self.normalize(queries)

        default = ann.search(queries, 5)
        ann.config["faiss"]["nprobe"] = ann.backend.nlist
        full = ann.search(queries, 5)

        self.assertEqual(default, full)

    def testFaissNprobeSmall(self):
        """
        Test Faiss default nprobe clamp on a small index
        """

        ann, _ = self.faissmodel(100, {"components": "IVF2,Flat"})
        self.assertEqual(ann.backend.nlist, 2)
        self.assertEqual(ann.nprobe(), ann.backend.nlist)

    def testFaissNprobeUnchanged(self):
        """
        Test Faiss default nprobe without a clamp
        """

        ann, _ = self.faissmodel(5001, {"components": "IVF16,Flat"})
        expected = round(ann.cells(ann.count()) / 16)
        self.assertLessEqual(expected, ann.backend.nlist)
        self.assertEqual(ann.nprobe(), expected)

    def testFaissQuantizeDisabled(self):
        """
        Test that quantize: false disables Faiss storage quantization
        """

        for quantize, expected in [(False, "IDMap,Flat"), (True, "IDMap,SQ8")]:
            with self.subTest(quantize=quantize):
                ann = ANNFactory.create({"backend": "faiss", "dimensions": 4, "quantize": quantize})
                self.assertEqual(ann.configure(100, 100), expected)

    def testGGML(self):
        """
        Test GGML backend
        """

        self.runTests("ggml")

    def testGGMLDelete(self):
        """
        Test GGML backend deletes
        """

        ann = ANNFactory.create({"backend": "ggml"})

        # Generate and index dummy data
        data = np.random.rand(100, 256).astype(np.float32)
        ann.index(data)

        # Validate count
        self.assertEqual(ann.count(), 100)

        # Test delete
        ann.delete([0])
        self.assertEqual(ann.count(), 99)

        # Deleting the same id again is a no-op for the count
        ann.delete([0])
        self.assertEqual(ann.count(), 99)

        # Save updated index with deletes and reload
        index = os.path.join(tempfile.gettempdir(), "ggml.deletes")
        ann.save(index)
        ann.load(index)
        self.assertEqual(ann.count(), 99)

    def testGGMLQuantizeDisabled(self):
        """
        Test that quantize: false disables GGML tensor quantization
        """

        data = np.random.rand(4, 256).astype(np.float32)
        for quantize, expected in [(False, ggml.GGML_TYPE_F32), (True, ggml.GGML_TYPE_Q8_0), (4, ggml.GGML_TYPE_Q4_0)]:
            with self.subTest(quantize=quantize):
                tensors = GGMLTensors(False, 64, quantize)
                self.assertEqual(tensors.tensortype(data), expected)

    def testGGMLEmpty(self):
        """
        Test GGML backend with an empty index
        """

        ann = ANNFactory.create({"backend": "ggml", "dimensions": 240})

        # Index an empty array
        ann.index(np.zeros((0, 240), dtype=np.float32))
        self.assertEqual(ann.count(), 0)

        # Test save and load
        index = os.path.join(tempfile.gettempdir(), "ggml.empty")
        ann.save(index)
        ann.load(index)
        self.assertEqual(ann.count(), 0)

        # Append data to the loaded empty index
        data = np.random.rand(10, 240).astype(np.float32)
        self.normalize(data)
        ann.append(data)
        self.assertEqual(ann.count(), 10)

        # Generate query vector and test search
        query = np.random.rand(240).astype(np.float32)
        self.normalize(query)
        self.assertGreater(ann.search(np.array([query]), 1)[0][0][1], 0)

    def testGGMLQuantization(self):
        """
        Test GGML backend with quantization enabled
        """

        ann = ANNFactory.create({"backend": "ggml", "ggml": {"quantize": "Q4_0"}})

        # Generate and index dummy data
        data = np.random.rand(100, 256).astype(np.float32)
        ann.index(data)

        # Test save and load
        index = os.path.join(tempfile.gettempdir(), "ggml.q4_0.v1")
        ann.save(index)
        ann.load(index)

        # Generate query vector and test search
        query = np.random.rand(256).astype(np.float32)
        self.normalize(query)
        self.assertGreater(ann.search(np.array([query]), 1)[0][0][1], 0)

        # Validate count
        self.assertEqual(ann.count(), 100)

        # Test delete
        ann.delete([0])
        self.assertEqual(ann.count(), 99)

        # Save updated index with deletes and reload
        index = os.path.join(tempfile.gettempdir(), "ggml.q4_0.v2")
        ann.save(index)
        ann.load(index)
        ann.index(data)

    def testGGMLInvalid(self):
        """
        Test invalid GGML configurations
        """

        data = np.random.rand(100, 240).astype(np.float32)

        with self.assertRaises(ValueError):
            ann = ANNFactory.create({"backend": "ggml", "ggml": {"quantize": "NOEXIST", "gpu": False}})
            ann.index(data)

        with self.assertRaises(ValueError):
            ann = ANNFactory.create({"backend": "ggml", "ggml": {"quantize": "Q4_K"}})
            ann.index(data)

    def testGGMLSingleRow(self):
        """
        Test GGML backend with a single row index
        """

        ann = ANNFactory.create({"backend": "ggml", "dimensions": 240})

        # Generate and index a single row
        data = np.random.rand(1, 240).astype(np.float32)
        self.normalize(data)
        ann.index(data)
        self.assertEqual(ann.count(), 1)

        # Test save and load
        index = os.path.join(tempfile.gettempdir(), "ggml.single")
        ann.save(index)
        ann.load(index)
        self.assertEqual(ann.count(), 1)
        self.assertEqual(ann.search(data, 1)[0][0][0], 0)

        # Append a row
        ann.append(data)
        self.assertEqual(ann.count(), 2)

        # Test delete, including an out of range id
        ann.delete([0, 5])
        self.assertEqual(ann.count(), 1)

    def testHnsw(self):
        """
        Test Hnswlib backend
        """

        self.runTests("hnsw")

    def testHnswCustom(self):
        """
        Test Hnswlib backend with custom settings
        """

        # Test with custom settings
        self.runTests("hnsw", {"hnsw": {"efconstruction": 100, "m": 8, "randomseed": 0, "efsearch": 15}})

    def testHnswLoad(self):
        """
        Test Hnswlib backend with a configuration that doesn't have a delete counter
        """

        # Save an index
        index = os.path.join(tempfile.gettempdir(), "hnsw.load")
        self.backend("hnsw", length=100).save(index)

        # Open the index with a configuration that only has the settings needed to load it
        ann = ANNFactory.create({"backend": "hnsw", "dimensions": 240, "metric": "ip"})
        ann.load(index)

        # Validate count, search and delete
        self.assertEqual(ann.count(), 100)
        self.assertEqual(len(ann.search(np.random.rand(1, 240).astype(np.float32), 1)[0]), 1)

        ann.delete([0])
        self.assertEqual(ann.count(), 99)

    def testMilvus(self):
        """
        Test milvus-lite backend
        """

        self.runTests("milvus")

    def testMilvusCustom(self):
        """
        Test milvus-lite backend with custom settings
        """

        self.runTests("milvus", {"milvus": {"m": 16}})

        # Test invalid file path handled
        with self.assertRaises(FileNotFoundError):
            ann = ANNFactory.create({"backend": "milvus"})
            ann.load("non-exist-path")

    def testNotImplemented(self):
        """
        Test exceptions for non-implemented methods
        """

        ann = ANN({})

        self.assertRaises(NotImplementedError, ann.load, None)
        self.assertRaises(NotImplementedError, ann.index, None)
        self.assertRaises(NotImplementedError, ann.append, None)
        self.assertRaises(NotImplementedError, ann.delete, None)
        self.assertRaises(NotImplementedError, ann.search, None, None)
        self.assertRaises(NotImplementedError, ann.count)
        self.assertRaises(NotImplementedError, ann.save, None)

    def testNumPy(self):
        """
        Test NumPy backend
        """

        self.runTests("numpy")

    @patch.dict(os.environ, {"ALLOW_PICKLE": "True"})
    def testNumPyLegacy(self):
        """
        Test NumPy backend with legacy pickled data
        """

        serializer = SerializeFactory.create("pickle", allowpickle=True)

        # Create output directory
        output = os.path.join(tempfile.gettempdir(), "ann.npy")
        path = os.path.join(output, "embeddings")
        os.makedirs(output, exist_ok=True)

        # Generate data and save as pickle
        data = np.random.rand(100, 240).astype(np.float32)
        serializer.save(data, path)

        ann = ANNFactory.create({"backend": "numpy"})
        ann.load(path)

        # Validate count
        self.assertEqual(ann.count(), 100)

    def testNumPySafetensors(self):
        """
        Test NumPy backend with safetensors storage
        """

        ann = ANNFactory.create({"backend": "numpy", "numpy": {"safetensors": True}})

        # Generate and index dummy data
        data = np.random.rand(100, 240).astype(np.float32)
        ann.index(data)

        # Test save and load
        index = os.path.join(tempfile.gettempdir(), "numpy.safetensors")
        ann.save(index)
        ann.load(index)

        # Generate query vector and test search
        query = np.random.rand(240).astype(np.float32)
        self.normalize(query)
        self.assertGreater(ann.search(np.array([query]), 1)[0][0][1], 0)

        # Validate count
        self.assertEqual(ann.count(), 100)

    def testNumPyQuantizeDelete(self):
        """
        Test that deleted rows don't resurface in search results for quantized (hamming) NumPy indexes
        """

        ann = ANNFactory.create({"backend": "numpy", "quantize": 1, "dimensions": 4})

        # Index uint8 vectors
        data = np.array([[1, 0, 0, 0], [0, 255, 0, 0], [0, 0, 255, 0], [0, 0, 0, 255]], dtype=np.uint8)
        ann.index(data)

        # Delete the first row and search with its (former) vector
        ann.delete([0])
        results = ann.search(np.array([data[0]]), 4)[0]

        # Deleted row must score 0 so the caller's score > 0 filter drops it, same as the dot-product path
        self.assertEqual(dict(results).get(0), 0)
        self.assertEqual(ann.count(), 3)

    @patch("sqlalchemy.orm.Query.limit")
    def testPGVector(self, query):
        """
        Test PGVector backend
        """

        # pylint: disable=W0613
        @compiles(BIT, "sqlite")
        def compile_bit_sqlite(type_, compiler, **kw):
            return "BLOB"

        # Generate test record
        data = np.random.rand(1, 240).astype(np.float32)

        # Mock database query
        query.return_value = [(x, -1.0) for x in range(data.shape[0])]

        configs = [
            ("full", {"dimensions": 240}, {}, data),
            ("half", {"dimensions": 240}, {"precision": "half"}, data),
            ("binary", {"quantize": 1, "dimensions": 240 * 8}, {}, data.astype(np.uint8)),
        ]

        # Create ANN
        for name, config, pgvector, data in configs:
            path = os.path.join(tempfile.gettempdir(), f"pgvector.{name}.sqlite")
            ann = ANNFactory.create(
                {**{"backend": "pgvector", "pgvector": {**{"url": f"sqlite:///{path}", "schema": "txtai"}, **pgvector}}, **config}
            )

            # Test indexing
            ann.index(data)
            ann.append(data)

            # Validate search results
            self.assertEqual(ann.search(data, 1), [[(0, 1.0)]])

            # Validate save/load/delete
            ann.save(None)
            ann.load(None)

            # Validate count
            self.assertEqual(ann.count(), 2)

            # Test delete
            ann.delete([0])
            self.assertEqual(ann.count(), 1)

            # Close ANN
            ann.close()

    def testRabitQ(self):
        """
        Test RabitQ backend
        """

        self.runTests("rabitq")

    def testRabitQCustom(self):
        """
        Test RabitQ backend with custom settings
        """

        # Test with custom settings
        self.runTests("rabitq", {"rabitq": {"mode": "hnsw"}}, False)
        self.runTests("rabitq", {"rabitq": {"clusters": 8, "nprobe": 2}}, False)
        self.runTests("rabitq", {"rabitq": {"nbits": 4}}, False)
        self.runTests("rabitq", {"rabitq": {"mode": "hnsw", "nbits": 4}}, False)
        self.runTests("rabitq", {"rabitq": {"nbits": 32}}, False)

        # Generate dummy data
        data = np.random.rand(100, 240).astype(np.float32)
        self.normalize(data)

        # Test invalid modes and quantization bits
        for mode, nbits in [("invalid", 1), ("ivf", 10), ("hnsw", 32)]:
            with self.assertRaises(ValueError):
                ann = ANNFactory.create({"backend": "rabitq", "dimensions": 240, "rabitq": {"mode": mode, "nbits": nbits}})
                ann.index(data)

        # Test a failed load leaves an empty index
        ann = ANNFactory.create({"backend": "rabitq", "dimensions": 240})
        ann.index(data)
        index = os.path.join(tempfile.gettempdir(), f"rabitq.invalid.{round(time.time() * 1000)}")
        ann.save(index)
        ann.config["rabitq"] = {"mode": "invalid"}
        with self.assertRaises(ValueError):
            ann.load(index)
        self.assertEqual(ann.count(), 0)

    def testRabitQDelete(self):
        """
        Test RabitQ deletes mark rows without rebuilding the index
        """

        # Generate dummy data
        data = np.random.rand(100, 240).astype(np.float32)
        self.normalize(data)

        for mode in ["ivf", "hnsw"]:
            # Probe all clusters so ivf searches can fill the limit
            ann = ANNFactory.create({"backend": "rabitq", "dimensions": 240, "rabitq": {"mode": mode, "nprobe": 100}})
            ann.index(data)
            backend = ann.backend

            # Empty deletes are ignored
            ann.delete([])
            self.assertEqual(ann.count(), 100)

            # Deleted rows are skipped by search and the limit is still filled
            ann.delete([0, 1])
            self.assertIs(ann.backend, backend)
            self.assertEqual(ann.count(), 98)
            for result in ann.search(data[:2], 10):
                self.assertEqual(len(result), 10)
                self.assertFalse({0, 1} & {uid for uid, _ in result})

            # Deletes are kept after a save and reload
            index = os.path.join(tempfile.gettempdir(), f"rabitq.{mode}.{round(time.time() * 1000)}")
            ann.save(index)
            ann.load(index)
            self.assertEqual(ann.count(), 98)
            ann.delete([2])
            self.assertEqual(ann.count(), 97)

            # Append rebuilds the index without the deleted rows
            ann.append(data[:1])
            self.assertEqual(ann.vectors.shape[0], 98)
            self.assertEqual(ann.count(), 98)

            # Rebuild once deleted rows outnumber live rows
            ann.delete(list(range(3, 60)))
            self.assertEqual(ann.vectors.shape[0], ann.count())
            self.assertEqual(ann.count(), 41)

            # Deleting all rows leaves an empty index, deletes on an empty index are ignored
            ann.delete(list(range(200)))
            ann.delete([0])
            self.assertEqual(ann.count(), 0)
            self.assertEqual(ann.search(data[:1], 10), [[]])

    @unittest.skipIf(platform.system() == "Darwin", "SQLite extensions not supported on macOS")
    def testSQLite(self):
        """
        Test SQLite backend
        """

        self.runTests("sqlite")

    @unittest.skipIf(platform.system() == "Darwin", "SQLite extensions not supported on macOS")
    def testSQLiteCustom(self):
        """
        Test SQLite backend with custom settings
        """

        # Test with custom settings
        self.runTests("sqlite", {"sqlite": {"quantize": 1}})
        self.runTests("sqlite", {"sqlite": {"quantize": 8}})

        # Test saving to a new path
        model = self.backend("sqlite")
        expected = model.count() - 1

        # Test save variations
        index = os.path.join(tempfile.gettempdir(), "ann.sqlite")
        new = os.path.join(tempfile.gettempdir(), "ann.sqlite.new")

        # Save new
        model.save(index)

        # Save to same path
        model.save(index)

        # Delete id
        model.delete([0])

        # Save to another path
        model.load(index)
        model.save(new)

        self.assertEqual(model.count(), expected)

    @unittest.skipIf(platform.system() == "Darwin", "SQLite extensions not supported on macOS")
    @unittest.skipIf(os.name == "nt", "SQLite copy skipped on Windows due to file locking")
    def testSQLiteSaveExistingPath(self):
        """
        Test saving a SQLite index to an existing path and overwriting the existing database
        """

        # Test saving to a new path
        model = self.backend("sqlite")

        # Test save variations
        index = os.path.join(tempfile.gettempdir(), "ann.sqlite.existing")

        # Save new
        model.save(index)

        # Test saving to a new path
        model = self.backend("sqlite")
        expected = model.count() - 1

        # Delete id
        model.delete([0])

        # Save to same path
        model.save(index)

        self.assertEqual(model.count(), expected)

    @unittest.skipIf(platform.system() == "Darwin", "SQLite extensions not supported on macOS")
    def testSQLiteSaveNewPath(self):
        """
        Test saving a loaded and modified SQLite index to a new path, then loading the new copy
        """

        for quantize in [None, 1, 8]:
            params = {"sqlite": {"quantize": quantize}} if quantize else None

            index = os.path.join(tempfile.gettempdir(), f"ann.sqlite.load.{quantize}")
            new = os.path.join(tempfile.gettempdir(), f"ann.sqlite.load.{quantize}.new")

            # Build and save index
            model = self.backend("sqlite", params, 500)
            model.save(index)
            model.close()

            # Load index, modify and save to a new path
            model = ANNFactory.create(model.config)
            model.load(index)
            model.delete([0])
            model.append(np.random.rand(10, 240).astype(np.float32))
            model.save(new)
            model.close()

            # Load new copy and check that it has all the changes
            model = ANNFactory.create(model.config)
            model.load(new)
            self.assertEqual(model.count(), 509)
            self.assertEqual(len(model.search(np.random.rand(1, 240).astype(np.float32), 10)[0]), 10)
            model.close()

    @unittest.skipIf(platform.system() == "Darwin", "SQLite extensions not supported on macOS")
    def testSQLiteQuantizeDisabled(self):
        """
        Test that quantize: false disables SQLite storage quantization
        """

        for quantize, expected in [(False, None), (True, 8), (1, 1)]:
            with self.subTest(quantize=quantize):
                ann = ANNFactory.create({"backend": "sqlite", "dimensions": 4, "sqlite": {"quantize": quantize}})
                self.assertEqual(ann.quantize, expected)

    def testTorch(self):
        """
        Test Torch backend
        """

        self.runTests("torch")

    def testTorchDelete(self):
        """
        Test Torch backend delete keeps the stored data type for quantized (uint8) and float16 indexes
        """

        for dtype, config in [(np.uint8, {"quantize": 1}), (np.float16, {})]:
            with self.subTest(dtype=dtype.__name__):
                ann = ANNFactory.create({"backend": "torch", "dimensions": 4, **config})

                # Index vectors with a single non-zero value per row
                data = np.array([[1, 0, 0, 0], [0, 255, 0, 0], [0, 0, 255, 0], [0, 0, 0, 255]], dtype=dtype)
                ann.index(data)

                # Delete the first row and search with its (former) vector
                ann.delete([0])
                results = ann.search(np.array([data[0]]), 4)[0]

                # Deleted row must score 0 so the caller's score > 0 filter drops it and the data type must be unchanged
                self.assertEqual(dict(results).get(0), 0)
                self.assertEqual(ann.count(), 3)
                self.assertEqual(str(ann.backend.dtype), f"torch.{dtype.__name__}")

    @unittest.skipIf(platform.system() == "Darwin", "Torch quantization not supported on macOS")
    def testTorchQuantization(self):
        """
        Test Torch backend with quantization enabled
        """

        for qtype in ["fp4", "nf4", "int8"]:
            ann = ANNFactory.create({"backend": "torch", "torch": {"quantize": {"type": qtype}}})

            # Generate and index dummy data
            data = np.random.rand(100, 240).astype(np.float32)
            ann.index(data)

            # Test save and load
            index = os.path.join(tempfile.gettempdir(), f"{qtype}.safetensors")
            ann.save(index)
            ann.load(index)

            # Generate query vector and test search
            query = np.random.rand(240).astype(np.float32)
            self.normalize(query)
            self.assertGreater(ann.search(np.array([query]), 1)[0][0][1], 0)

            # Validate count
            self.assertEqual(ann.count(), 100)

            # Test delete
            ann.delete([0])
            self.assertEqual(ann.count(), 99)

    def testTurboVec(self):
        """
        Test turbovec backend
        """

        self.runTests("turbovec")

    def testZvec(self):
        """
        Test zvec backend
        """

        self.runTests("zvec")

    def testZvecCustom(self):
        """
        Test zvec backend with custom settings
        """

        self.runTests("zvec", {"zvec": {"m": 16}})

        # Test invalid file path handled
        with self.assertRaises(FileNotFoundError):
            ann = ANNFactory.create({"backend": "zvec"})
            ann.load("non-exist-path")

    def runTests(self, name, params=None, update=True):
        """
        Runs a series of standard backend tests.

        Args:
            name: backend name
            params: additional config parameters
            update: If append/delete options should be tested
        """

        self.assertEqual(self.backend(name, params).config["backend"], name)
        self.assertEqual(self.save(name, params).count(), 10000)

        if update:
            self.assertEqual(self.append(name, params, 500).count(), 10500)
            self.assertEqual(self.delete(name, params, [0, 1]).count(), 9998)
            self.assertEqual(self.delete(name, params, [100000]).count(), 10000)

        self.assertGreater(self.search(name, params), 0)
        self.assertGreater(self.limit(name, params), 0)

    def backend(self, name, params=None, length=10000):
        """
        Test a backend.

        Args:
            name: backend name
            params: additional config parameters
            length: number of rows to generate

        Returns:
            ANN model
        """

        # Generate test data
        data = np.random.rand(length, 240).astype(np.float32)
        self.normalize(data)

        config = {"backend": name, "dimensions": data.shape[1]}
        if params:
            config.update(params)

        model = ANNFactory.create(config)
        model.index(data)

        return model

    def faissmodel(self, length, settings):
        """
        Builds a seeded Faiss model.

        Args:
            length: number of rows to generate
            settings: Faiss settings

        Returns:
            Faiss model and random number generator
        """

        rng = np.random.default_rng(0)
        data = rng.random((length, 8), dtype=np.float32)
        self.normalize(data)

        model = ANNFactory.create({"backend": "faiss", "dimensions": data.shape[1], "faiss": settings})
        model.index(data)

        return model, rng

    def append(self, name, params=None, length=500):
        """
        Appends new data to index.

        Args:
            name: backend name
            params: additional config parameters
            length: number of rows to generate

        Returns:
            ANN model
        """

        # Initial model
        model = self.backend(name, params)

        # Generate test data
        data = np.random.rand(length, 240).astype(np.float32)
        self.normalize(data)

        model.append(data)

        return model

    def delete(self, name, params=None, ids=None):
        """
        Deletes data from index.

        Args:
            name: backend name
            params: additional config parameters
            ids: ids to delete

        Returns:
            ANN model
        """

        # Initial model
        model = self.backend(name, params)
        model.delete(ids)

        return model

    def save(self, name, params=None):
        """
        Test save/load.

        Args:
            name: backend name
            params: additional config parameters

        Returns:
            ANN model
        """

        model = self.backend(name, params)

        # Generate temp file path
        index = os.path.join(tempfile.gettempdir(), f"ann.{name}.{round(time.time() * 1000)}")

        # Generate query vector
        query = np.random.rand(240).astype(np.float32)
        self.normalize(query)

        # Save index and ensure it's still searchable
        model.save(index)
        self.assertGreater(model.search(np.array([query]), 1)[0][0][1], 0)

        # Close and reload index
        model.close()
        model.load(index)

        # Ensure reloaded index is searchable
        self.assertGreater(model.search(np.array([query]), 1)[0][0][1], 0)

        return model

    def search(self, name, params=None):
        """
        Test ANN search.

        Args:
            name: backend name
            params: additional config parameters

        Returns:
            search results
        """

        # Generate ANN index
        model = self.backend(name, params)

        # Generate query vector
        query = np.random.rand(240).astype(np.float32)
        self.normalize(query)

        # Ensure top result has similarity > 0
        return model.search(np.array([query]), 1)[0][0][1]

    def limit(self, name, params=None):
        """
        Test ANN limit search.

        Args:
            name: backend name
            params: additional config parameters

        Returns:
            search results
        """

        # Generate ANN index
        model = self.backend(name, params, 50)

        # Generate query vector
        query = np.random.rand(240).astype(np.float32)
        self.normalize(query)

        # Ensure limit being > count doesn't throw an error
        return model.search(np.array([query]), 100)[0][0][1]

    def normalize(self, embeddings):
        """
        Normalizes embeddings using L2 normalization. Operation applied directly on array.

        Args:
            embeddings: input embeddings matrix
        """

        # Calculation is different for matrices vs vectors
        if len(embeddings.shape) > 1:
            embeddings /= np.linalg.norm(embeddings, axis=1)[:, np.newaxis]
        else:
            embeddings /= np.linalg.norm(embeddings)