# 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.

"""
Comprehensive test suite for IndexIVFFlatPanorama.

Panorama is an adaptation of IndexIVFFlat that uses level-oriented storage
and progressive filtering with Cauchy-Schwarz bounds to achieve significant
speedups (up to 20x when combined with PCA or Cayley transforms) with zero
loss in accuracy.

Paper: https://www.arxiv.org/pdf/2510.00566
"""

import unittest

import faiss
import numpy as np
from faiss.contrib.datasets import SyntheticDataset


class TestIndexIVFFlatPanorama(unittest.TestCase):
    """Test Suite for IndexIVFFlatPanorama."""

    # Metrics to test
    METRICS = [faiss.METRIC_L2, faiss.METRIC_INNER_PRODUCT]

    # Helper methods for index creation and data generation

    def generate_data(self, d, nt, nb, nq, seed=42):
        ds = SyntheticDataset(d, nt, nb, nq, seed=seed)
        return ds.get_train(), ds.get_database(), ds.get_queries()

    def create_ivf_flat(
        self,
        d,
        nlist,
        xt,
        xb=None,
        nprobe=None,
        make_direct_map=False,
        metric=faiss.METRIC_L2,
    ):
        quantizer = (
            faiss.IndexFlatL2(d) if metric == faiss.METRIC_L2 else faiss.IndexFlatIP(d)
        )
        index = faiss.IndexIVFFlat(quantizer, d, nlist, metric)
        index.train(xt)
        if make_direct_map:
            index.make_direct_map()
        if xb is not None:
            index.add(xb)
        if nprobe is not None:
            index.nprobe = nprobe
        return index

    def create_panorama(
        self,
        d,
        nlist,
        nlevels,
        xt,
        xb=None,
        nprobe=None,
        make_direct_map=False,
        metric=faiss.METRIC_L2,
    ):
        """Create and initialize IndexIVFFlatPanorama."""
        quantizer = (
            faiss.IndexFlatL2(d) if metric == faiss.METRIC_L2 else faiss.IndexFlatIP(d)
        )
        index = faiss.IndexIVFFlatPanorama(quantizer, d, nlist, nlevels, metric)
        index.train(xt)
        if make_direct_map:
            index.make_direct_map()
        if xb is not None:
            index.add(xb)
        if nprobe is not None:
            index.nprobe = nprobe
        return index

    def assert_search_results_equal(
        self,
        D_regular,
        I_regular,
        D_panorama,
        I_panorama,
        rtol=1e-5,
        atol=1e-7,
        otol=1e-3,
    ):
        # Allow small tolerance in overlap rate to account for floating-point errors
        # in distance computations that can affect ordering when distances are nearly equal.
        # Faiss: (a - b) * (a - b) vs. Panorama: a * a + b * b - 2(a * b)
        overlap_rate = np.mean(I_regular == I_panorama)

        self.assertGreater(
            overlap_rate,
            1 - otol,
            f"Overlap rate {overlap_rate:.6f} is not > {1-otol:.3f}. ",
        )
        np.testing.assert_allclose(
            D_regular, D_panorama, rtol=rtol, atol=atol, err_msg="Distances mismatch"
        )

    def assert_range_results_equal(
        self,
        lims_regular,
        D_regular,
        I_regular,
        lims_panorama,
        D_panorama,
        I_panorama,
        nq,
        otol=1e-3,
        rtol=1e-4,
    ):
        """Compare range search results with tolerance for boundary differences."""
        total_matches = total_regular = 0

        for i in range(nq):
            ids_reg = I_regular[lims_regular[i] : lims_regular[i + 1]]
            ids_pan = I_panorama[lims_panorama[i] : lims_panorama[i + 1]]
            total_regular += len(ids_reg)

            common = np.intersect1d(ids_reg, ids_pan)
            total_matches += len(common)

            if len(common) > 0:
                # Extract and sort distances for common IDs
                mask_reg, mask_pan = np.isin(ids_reg, common), np.isin(ids_pan, common)
                dist_reg = D_regular[lims_regular[i] : lims_regular[i + 1]]
                dist_pan = D_panorama[lims_panorama[i] : lims_panorama[i + 1]]

                sort_reg, sort_pan = np.argsort(ids_reg[mask_reg]), np.argsort(
                    ids_pan[mask_pan]
                )
                np.testing.assert_allclose(
                    dist_reg[mask_reg][sort_reg],
                    dist_pan[mask_pan][sort_pan],
                    rtol=rtol,
                    err_msg=f"Query {i}: distance mismatch",
                )

        if total_regular > 0:
            overlap_rate = total_matches / total_regular
            self.assertGreater(
                overlap_rate,
                1 - otol,
                f"Range overlap {overlap_rate:.6f} not > {1-otol:.3f}",
            )

    def validate_and_compare_range_results(
        self, metric, radius, lims_reg, D_reg, I_reg, lims_pan, D_pan, I_pan, nq
    ):
        """Helper to validate range search results match between regular and panorama."""
        if metric == faiss.METRIC_L2:
            self.assertTrue(
                np.all(D_pan <= radius), f"Some L2 distances exceed radius={radius}"
            )
        else:
            self.assertTrue(
                np.all(D_pan >= radius), f"Some IP distances below radius={radius}"
            )
        self.assert_range_results_equal(
            lims_reg, D_reg, I_reg, lims_pan, D_pan, I_pan, nq
        )

    # Core functionality tests

    def test_exact_match_with_ivf_flat(self):
        """Core test: Panorama must return identical results to IndexIVFFlat"""
        d, nb, nt, nq, nlist, nlevels, k = 128, 100000, 150000, 1000, 512, 8, 20
        xt, xb, xq = self.generate_data(d, nt, nb, nq, seed=42)

        for metric in self.METRICS:
            for nprobe in [1, 16, 64, 256]:
                with self.subTest(metric=metric, nprobe=nprobe):
                    index_regular = self.create_ivf_flat(
                        d, nlist, xt, xb, nprobe, metric=metric
                    )
                    index_panorama = self.create_panorama(
                        d, nlist, nlevels, xt, xb, nprobe, metric=metric
                    )

                    D_regular, I_regular = index_regular.search(xq, k)
                    D_panorama, I_panorama = index_panorama.search(xq, k)

                    self.assert_search_results_equal(
                        D_regular, I_regular, D_panorama, I_panorama
                    )

    def test_exact_match_with_ivf_flat_medium(self):
        """Core test: Medium scale version"""
        d, nb, nt, nq, nlist, nlevels, k = 64, 10000, 15000, 200, 128, 4, 10
        xt, xb, xq = self.generate_data(d, nt, nb, nq, seed=42)

        for metric in self.METRICS:
            for nprobe in [1, 4, 8, nlist]:
                with self.subTest(metric=metric, nprobe=nprobe):
                    index_regular = self.create_ivf_flat(
                        d, nlist, xt, xb, nprobe, metric=metric
                    )
                    index_panorama = self.create_panorama(
                        d, nlist, nlevels, xt, xb, nprobe, metric=metric
                    )

                    D_regular, I_regular = index_regular.search(xq, k)
                    D_panorama, I_panorama = index_panorama.search(xq, k)

                    self.assert_search_results_equal(
                        D_regular, I_regular, D_panorama, I_panorama
                    )

    def test_range_search(self):
        """Test range search returns correct results within radius"""
        d, nb, nt, nq, nlist, nlevels = 128, 50000, 75000, 500, 256, 8
        xt, xb, xq = self.generate_data(d, nt, nb, nq, seed=123)

        for metric in self.METRICS:
            index_regular = self.create_ivf_flat(
                d, nlist, xt, xb, nprobe=64, metric=metric
            )
            index_panorama = self.create_panorama(
                d, nlist, nlevels, xt, xb, nprobe=64, metric=metric
            )

            for radius in [0.5, 1.0, 2.0, 5.0]:
                with self.subTest(metric=metric, radius=radius):
                    lims_reg, D_reg, I_reg = index_regular.range_search(xq, radius)
                    lims_pan, D_pan, I_pan = index_panorama.range_search(xq, radius)
                    self.validate_and_compare_range_results(
                        metric,
                        radius,
                        lims_reg,
                        D_reg,
                        I_reg,
                        lims_pan,
                        D_pan,
                        I_pan,
                        nq,
                    )

    # Parameter variation tests

    def test_different_n_levels(self):
        """Test correctness with various n_levels parameter values"""
        d, nb, nt, nq, nlist, k = 128, 25000, 40000, 300, 128, 15
        xt, xb, xq = self.generate_data(d, nt, nb, nq, seed=456)

        for metric in self.METRICS:
            index_base = self.create_ivf_flat(
                d, nlist, xt, xb, nprobe=32, metric=metric
            )
            D_base, I_base = index_base.search(xq, k)

            nt = faiss.omp_get_max_threads()
            faiss.omp_set_num_threads(1)

            prev_ratio_dims_scanned = float("inf")
            for nlevels in [1, 2, 4, 8, 16, 32]:
                with self.subTest(metric=metric, nlevels=nlevels):
                    faiss.cvar.indexPanorama_stats.reset()
                    index = self.create_panorama(
                        d, nlist, nlevels, xt, xb, nprobe=32, metric=metric
                    )
                    D, I = index.search(xq, k)
                    self.assert_search_results_equal(D_base, I_base, D, I)

            # Verify ratio_dims_scanned decreases as n_levels increases
            ratio_dims_scanned = faiss.cvar.indexPanorama_stats.ratio_dims_scanned
            self.assertLess(ratio_dims_scanned, prev_ratio_dims_scanned)
            prev_ratio_dims_scanned = ratio_dims_scanned

            faiss.omp_set_num_threads(nt)

    def test_uneven_dimension_division(self):
        """Test when n_levels doesn't evenly divide dimension"""
        test_cases = [(65, 4), (63, 8), (100, 7)]

        # TODO(aknayar): Test functions like get_single_code().

        for metric in self.METRICS:
            for d, nlevels in test_cases:
                with self.subTest(metric=metric, d=d, nlevels=nlevels):
                    nb, nt, nq, nlist, k = 5000, 700, 20, 16, 5
                    xt, xb, xq = self.generate_data(d, nt, nb, nq, seed=789)

                    index_regular = self.create_ivf_flat(
                        d, nlist, xt, xb, nprobe=4, metric=metric
                    )
                    index_panorama = self.create_panorama(
                        d, nlist, nlevels, xt, xb, nprobe=4, metric=metric
                    )

                    D_regular, I_regular = index_regular.search(xq, k)
                    D_panorama, I_panorama = index_panorama.search(xq, k)

                    self.assert_search_results_equal(
                        D_regular, I_regular, D_panorama, I_panorama
                    )

    def test_single_level(self):
        """Test edge case with n_levels=1"""
        d, nb, nt, nq, nlist, nlevels, k = 32, 500, 700, 20, 16, 1, 5
        xt, xb, xq = self.generate_data(d, nt, nb, nq, seed=333)

        for metric in self.METRICS:
            with self.subTest(metric=metric):
                index_regular = self.create_ivf_flat(
                    d, nlist, xt, xb, nprobe=4, metric=metric
                )
                index_panorama = self.create_panorama(
                    d, nlist, nlevels, xt, xb, nprobe=4, metric=metric
                )

                D_regular, I_regular = index_regular.search(xq, k)
                D_panorama, I_panorama = index_panorama.search(xq, k)

                self.assert_search_results_equal(
                    D_regular, I_regular, D_panorama, I_panorama
                )

    def test_multiple_levels_small_dimension(self):
        """Test edge case: more levels than dimension naturally supports"""
        d, nb, nt, nq, nlist, nlevels, k = 16, 200, 300, 10, 4, 8, 5
        xt, xb, xq = self.generate_data(d, nt, nb, nq, seed=1212)

        for metric in self.METRICS:
            with self.subTest(metric=metric):
                index_regular = self.create_ivf_flat(
                    d, nlist, xt, xb, nprobe=4, metric=metric
                )
                index_panorama = self.create_panorama(
                    d, nlist, nlevels, xt, xb, nprobe=4, metric=metric
                )

                D_regular, I_regular = index_regular.search(xq, k)
                D_panorama, I_panorama = index_panorama.search(xq, k)

                self.assert_search_results_equal(
                    D_regular, I_regular, D_panorama, I_panorama
                )

    # ID selector tests

    def test_id_selector_range(self):
        """Test ID filtering with range selector"""
        d, nb, nt, nq, nlist, nlevels, k = 128, 80000, 120000, 500, 512, 8, 20
        xt, xb, xq = self.generate_data(d, nt, nb, nq, seed=321)

        for metric in self.METRICS:
            with self.subTest(metric=metric):
                index_regular = self.create_ivf_flat(
                    d, nlist, xt, xb, nprobe=64, metric=metric
                )
                index_panorama = self.create_panorama(
                    d, nlist, nlevels, xt, xb, nprobe=64, metric=metric
                )

                params = faiss.SearchParametersIVF()
                params.sel = faiss.IDSelectorRange(10000, 50000)

                D_regular, I_regular = index_regular.search(xq, k, params=params)
                D_panorama, I_panorama = index_panorama.search(xq, k, params=params)

                self.assertTrue(np.all(I_panorama >= 10000))
                self.assertTrue(np.all(I_panorama < 50000))

                np.testing.assert_array_equal(I_regular, I_panorama)
                np.testing.assert_allclose(D_regular, D_panorama, rtol=1e-4)

    def test_id_selector_batch(self):
        """Test ID filtering with batch selector"""
        d, nb, nt, nq, nlist, nlevels, k = 128, 60000, 90000, 400, 256, 8, 20
        xt, xb, xq = self.generate_data(d, nt, nb, nq, seed=654)

        for metric in self.METRICS:
            with self.subTest(metric=metric):
                index_regular = self.create_ivf_flat(
                    d, nlist, xt, xb, nprobe=32, metric=metric
                )
                index_panorama = self.create_panorama(
                    d, nlist, nlevels, xt, xb, nprobe=32, metric=metric
                )

                allowed_ids = np.array([i * 100 for i in range(500)], dtype=np.int64)
                params = faiss.SearchParametersIVF()
                params.sel = faiss.IDSelectorBatch(allowed_ids)

                D_regular, I_regular = index_regular.search(xq, k, params=params)
                D_panorama, I_panorama = index_panorama.search(xq, k, params=params)

                allowed_set = set(allowed_ids) | set([-1])
                for id_val in I_panorama.flatten():
                    self.assertIn(int(id_val), allowed_set)

                np.testing.assert_array_equal(I_regular, I_panorama)
                np.testing.assert_allclose(D_regular, D_panorama, rtol=1e-4)

    def test_selector_with_small_dataset(self):
        """Test ID selectors with dataset smaller than batch size"""
        d, nb, nt, nq, nlist, nlevels, k = 32, 100, 200, 10, 4, 4, 5
        xt, xb, xq = self.generate_data(d, nt, nb, nq, seed=888)

        for metric in self.METRICS:
            with self.subTest(metric=metric):
                index_regular = self.create_ivf_flat(
                    d, nlist, xt, xb, nprobe=nlist, metric=metric
                )
                index_panorama = self.create_panorama(
                    d, nlist, nlevels, xt, xb, nprobe=nlist, metric=metric
                )

                params = faiss.SearchParametersIVF()
                params.sel = faiss.IDSelectorRange(20, 60)

                D_regular, I_regular = index_regular.search(xq, k, params=params)
                D_panorama, I_panorama = index_panorama.search(xq, k, params=params)

                self.assertTrue(np.all(I_panorama >= 20))
                self.assertTrue(np.all(I_panorama < 60))

                np.testing.assert_array_equal(I_regular, I_panorama)
                np.testing.assert_allclose(D_regular, D_panorama, rtol=1e-4)

    def test_selector_excludes_all(self):
        """Test selector that excludes all results"""
        d, nb, nt, nq, nlist, nlevels, k = 32, 300, 400, 5, 8, 4, 10
        xt, xb, xq = self.generate_data(d, nt, nb, nq, seed=999)

        for metric in self.METRICS:
            with self.subTest(metric=metric):
                index_panorama = self.create_panorama(
                    d, nlist, nlevels, xt, xb, nprobe=nlist, metric=metric
                )

                params = faiss.SearchParametersIVF()
                params.sel = faiss.IDSelectorRange(nb + 100, nb + 200)

                D, I = index_panorama.search(xq, k, params=params)
                self.assertTrue(np.all(I == -1))

    # Batch size and edge case tests

    def test_batch_boundaries(self):
        """Test correctness at various batch size boundaries (kBatchSize=256)"""
        d, nlist, nlevels, nt, nq, k = 128, 64, 8, 10000, 200, 15
        xt = np.random.rand(nt, d).astype("float32")
        xq = np.random.rand(nq, d).astype("float32")

        for metric in self.METRICS:
            for nb in [5000, 10000, 20000, 50000, 75000]:
                with self.subTest(metric=metric, nb=nb):
                    np.random.seed(987)
                    xb = np.random.rand(nb, d).astype("float32")

                    index_regular = self.create_ivf_flat(
                        d, nlist, xt, xb, nprobe=16, metric=metric
                    )
                    index_panorama = self.create_panorama(
                        d, nlist, nlevels, xt, xb, nprobe=16, metric=metric
                    )

                    D_regular, I_regular = index_regular.search(xq, k)
                    D_panorama, I_panorama = index_panorama.search(xq, k)

                    self.assert_search_results_equal(
                        D_regular, I_regular, D_panorama, I_panorama
                    )

    def test_empty_result_handling(self):
        """Test handling of empty search results"""
        d, nb, nt, nq, nlist, nlevels, k = 32, 100, 200, 5, 8, 4, 10
        xt, xb, _ = self.generate_data(d, nt, nb, nq, seed=111)
        xq = np.random.rand(nq, d).astype("float32") + 10.0  # Queries far from database

        for metric in self.METRICS:
            with self.subTest(metric=metric):
                index = self.create_panorama(
                    d, nlist, nlevels, xt, xb, nprobe=1, metric=metric
                )
                D, I = index.search(xq, k)

                self.assertEqual(D.shape, (nq, k))
                self.assertEqual(I.shape, (nq, k))

    def test_very_small_dataset(self):
        """Test with dataset smaller than batch size (< 128 vectors)"""
        test_cases = [10, 50, 100]

        for metric in self.METRICS:
            for nb in test_cases:
                with self.subTest(metric=metric, nb=nb):
                    d, nt, nlist, nlevels, nq = 32, max(nb, 100), 4, 4, 5
                    k = min(3, nb)
                    xt, xb, xq = self.generate_data(d, nt, nb, nq, seed=666 + nb)

                    index_regular = self.create_ivf_flat(
                        d, nlist, xt, xb, nprobe=nlist, metric=metric
                    )
                    index_panorama = self.create_panorama(
                        d, nlist, nlevels, xt, xb, nprobe=nlist, metric=metric
                    )

                    D_regular, I_regular = index_regular.search(xq, k)
                    D_panorama, I_panorama = index_panorama.search(xq, k)

                    self.assert_search_results_equal(
                        D_regular, I_regular, D_panorama, I_panorama
                    )

    def test_single_vector_per_cluster(self):
        """Test extreme case where clusters have very few vectors"""
        d, nb, nt, nq, nlist, nlevels, k = 32, 20, 100, 5, 16, 4, 3
        xt, xb, xq = self.generate_data(d, nt, nb, nq, seed=1313)

        for metric in self.METRICS:
            with self.subTest(metric=metric):
                quantizer1 = (
                    faiss.IndexFlatL2(d)
                    if metric == faiss.METRIC_L2
                    else faiss.IndexFlatIP(d)
                )
                quantizer2 = (
                    faiss.IndexFlatL2(d)
                    if metric == faiss.METRIC_L2
                    else faiss.IndexFlatIP(d)
                )
                index_regular = faiss.IndexIVFFlat(quantizer1, d, nlist, metric)
                index_panorama = faiss.IndexIVFFlatPanorama(
                    quantizer2, d, nlist, nlevels, metric
                )

                index_regular.cp.min_points_per_centroid = 1
                index_panorama.cp.min_points_per_centroid = 1

                index_regular.train(xt)
                index_regular.add(xb)
                index_regular.nprobe = nlist

                index_panorama.train(xt)
                index_panorama.add(xb)
                index_panorama.nprobe = nlist

                D_regular, I_regular = index_regular.search(xq, k)
                D_panorama, I_panorama = index_panorama.search(xq, k)

                self.assert_search_results_equal(
                    D_regular, I_regular, D_panorama, I_panorama
                )

    def test_range_search_edge_cases(self):
        """Test range search with extreme radius values"""
        d, nb, nt, nq, nlist, nlevels = 32, 500, 700, 10, 8, 4
        xt, xb, xq = self.generate_data(d, nt, nb, nq, seed=777)

        for metric in self.METRICS:
            index_regular = self.create_ivf_flat(
                d, nlist, xt, xb, nprobe=nlist, metric=metric
            )
            index_panorama = self.create_panorama(
                d, nlist, nlevels, xt, xb, nprobe=nlist, metric=metric
            )

            for radius in [0.01, 0.1, 100.0, 1000.0]:
                with self.subTest(metric=metric, radius=radius):
                    lims_reg, D_reg, I_reg = index_regular.range_search(xq, radius)
                    lims_pan, D_pan, I_pan = index_panorama.range_search(xq, radius)
                    self.validate_and_compare_range_results(
                        metric,
                        radius,
                        lims_reg,
                        D_reg,
                        I_reg,
                        lims_pan,
                        D_pan,
                        I_pan,
                        nq,
                    )

    # Dynamic operations tests

    def test_incremental_add(self):
        """Test adding vectors incrementally in multiple batches"""
        d, nt, nlist, nlevels, k = 128, 20000, 256, 8, 15
        xt = np.random.rand(nt, d).astype("float32")

        for metric in self.METRICS:
            with self.subTest(metric=metric):
                index_regular = self.create_ivf_flat(
                    d, nlist, xt, nprobe=32, metric=metric
                )
                index_panorama = self.create_panorama(
                    d, nlist, nlevels, xt, nprobe=32, metric=metric
                )

                batch_sizes = [10000, 15000, 20000, 5000]
                for batch_size in batch_sizes:
                    xb_batch = np.random.rand(batch_size, d).astype("float32")
                    index_regular.add(xb_batch)
                    index_panorama.add(xb_batch)

                nq = 300
                xq = np.random.rand(nq, d).astype("float32")

                D_regular, I_regular = index_regular.search(xq, k)
                D_panorama, I_panorama = index_panorama.search(xq, k)

                self.assert_search_results_equal(
                    D_regular, I_regular, D_panorama, I_panorama
                )

    def test_add_search_add_search(self):
        """Test interleaved add and search operations"""
        d, nt, nlist, nlevels, k = 32, 500, 8, 4, 5
        np.random.seed(555)
        xt = np.random.rand(nt, d).astype("float32")

        for metric in self.METRICS:
            with self.subTest(metric=metric):
                index_regular = self.create_ivf_flat(d, nlist, xt, metric=metric)
                index_panorama = self.create_panorama(
                    d, nlist, nlevels, xt, metric=metric
                )

                # First add and search
                xb1 = np.random.rand(200, d).astype("float32")
                index_regular.add(xb1)
                index_panorama.add(xb1)

                xq1 = np.random.rand(10, d).astype("float32")
                index_regular.nprobe = 4
                index_panorama.nprobe = 4

                D_reg_1, I_reg_1 = index_regular.search(xq1, k)
                D_pan_1, I_pan_1 = index_panorama.search(xq1, k)
                self.assert_search_results_equal(D_reg_1, I_reg_1, D_pan_1, I_pan_1)

                # Second add and search
                xb2 = np.random.rand(300, d).astype("float32")
                index_regular.add(xb2)
                index_panorama.add(xb2)

                xq2 = np.random.rand(10, d).astype("float32")
                D_reg_2, I_reg_2 = index_regular.search(xq2, k)
                D_pan_2, I_pan_2 = index_panorama.search(xq2, k)
                self.assert_search_results_equal(D_reg_2, I_reg_2, D_pan_2, I_pan_2)

    def test_update_vectors(self):
        """Test update operations (single, batch, and interleaved with search)"""
        d, nb, nt, nq, nlist, nlevels, k = 128, 40000, 60000, 400, 256, 8, 15
        xt, xb, xq = self.generate_data(d, nt, nb, nq, seed=1414)

        for metric in self.METRICS:
            with self.subTest(metric=metric):
                index_regular = self.create_ivf_flat(
                    d, nlist, xt, xb, nprobe=32, make_direct_map=True, metric=metric
                )
                index_panorama = self.create_panorama(
                    d,
                    nlist,
                    nlevels,
                    xt,
                    xb,
                    nprobe=32,
                    make_direct_map=True,
                    metric=metric,
                )

                # Test batch update
                update_ids = np.array(
                    [100, 1000, 5000, 10000, 20000, 30000], dtype=np.int64
                )
                xb_new = np.random.rand(len(update_ids), d).astype("float32")
                index_regular.update_vectors(update_ids, xb_new)
                index_panorama.update_vectors(update_ids, xb_new)

                D_reg, I_reg = index_regular.search(xq, k)
                D_pan, I_pan = index_panorama.search(xq, k)
                self.assert_search_results_equal(D_reg, I_reg, D_pan, I_pan)

                # Test interleaved update/search
                update_ids_2 = np.array([50, 500, 2500, 15000], dtype=np.int64)
                xb_new_2 = np.random.rand(len(update_ids_2), d).astype("float32")
                index_regular.update_vectors(update_ids_2, xb_new_2)
                index_panorama.update_vectors(update_ids_2, xb_new_2)

                D_reg_2, I_reg_2 = index_regular.search(xq, k)
                D_pan_2, I_pan_2 = index_panorama.search(xq, k)
                self.assert_search_results_equal(D_reg_2, I_reg_2, D_pan_2, I_pan_2)

    def test_serialization(self):
        """Test that writing and reading Panorama indexes preserves search results"""
        d, nb, nt, nq, nlist, nlevels, k = 128, 10000, 15000, 100, 128, 8, 20
        xt, xb, xq = self.generate_data(d, nt, nb, nq, seed=2024)

        for metric in self.METRICS:
            with self.subTest(metric=metric):
                index = self.create_panorama(
                    d, nlist, nlevels, xt, xb, nprobe=32, metric=metric
                )

                D_before, I_before = index.search(xq, k)
                index_after = faiss.deserialize_index(faiss.serialize_index(index))
                D_after, I_after = index_after.search(xq, k)

                np.testing.assert_array_equal(I_before, I_after)
                np.testing.assert_array_equal(D_before, D_after)

    def test_ratio_dims_scanned(self):
        """Test the correctness of the ratio of dimensions scanned"""
        d, nb, nq, nlist, k = 128, 500000, 1, 1, 1

        # Setup: All vectors in the dataset are [1, 1, ..., 1], except for
        # one, which is [0, 0, ..., 0]. The query is also [0, 0, ..., 0]. This
        # ensures that pruning happens as early as it can. As a result, the
        # ratio of dimensions scanned should be approximately (1 / nlevels).
        xb = np.ones((nb, d)).astype("float32")
        xb[-1] = 0
        xt = xb
        xq = np.zeros((nq, d)).astype("float32")

        for metric in self.METRICS:
            with self.subTest(metric=metric):
                index_base = self.create_ivf_flat(
                    d, nlist, xt, xb, nprobe=1, metric=metric
                )
                D_base, I_base = index_base.search(xq, k)

                nt = faiss.omp_get_max_threads()
                faiss.omp_set_num_threads(1)

                ratios = []
                nlevels_list = [1, 2, 4, 8, 16, 32]
                for nlevels in nlevels_list:
                    with self.subTest(nlevels=nlevels):
                        faiss.cvar.indexPanorama_stats.reset()
                        index = self.create_panorama(
                            d, nlist, nlevels, xt, xb, nprobe=1, metric=metric
                        )
                        D, I = index.search(xq, k)
                        self.assert_search_results_equal(D_base, I_base, D, I)

                        ratios.append(faiss.cvar.indexPanorama_stats.ratio_dims_scanned)

                expected_ratios = [1 / nlevels for nlevels in nlevels_list]
                np.testing.assert_allclose(ratios, expected_ratios, atol=1e-3)

                faiss.omp_set_num_threads(nt)
