"""Pre-filter vs post-filter vs in-graph filter for vector search with metadata.

Synthetic clustered, normalized float32 vectors. Each vector has a tenant id.
Two layouts: "random" (tenant independent of content) and "topical" (tenant
owns whole topic clusters, so a tenant's queries land in its own region).
Ground truth = exact search restricted to the tenant's vectors.
"""
import time, platform, numpy as np, hnswlib

N, D, K, NQ, CLUSTERS = 100_000, 128, 10, 200, 1000
rng = np.random.default_rng(0)
centers = rng.normal(size=(CLUSTERS, D)).astype(np.float32)
cl = rng.integers(0, CLUSTERS, N)
X = centers[cl] + 0.6 * rng.normal(size=(N, D)).astype(np.float32)
X /= np.linalg.norm(X, axis=1, keepdims=True)

index = hnswlib.Index(space="ip", dim=D)
index.init_index(max_elements=N, ef_construction=200, M=16)
index.add_items(X, np.arange(N))
index.set_ef(100)
index.set_num_threads(1)

def exact(allowed_mask, q):
    ids = np.flatnonzero(allowed_mask)
    s = X[ids] @ q
    top = np.argpartition(-s, K)[:K]
    return ids[top[np.argsort(-s[top])]]

def timed(fn, qs):
    out, ts = [], []
    for q in qs:
        t = time.perf_counter(); out.append(fn(q)); ts.append((time.perf_counter() - t) * 1000)
    return out, float(np.median(ts))

def recall(res, truth):
    return float(np.mean([len(set(r) & set(t)) / K for r, t in zip(res, truth)]))

def fill(res):
    return float(np.mean([len(r) for r in res]))

print(f"N={N} dim={D} k={K} queries/setting={NQ} hnsw M=16 efC=200 ef=100 threads=1")
print(f"python {platform.python_version()} numpy {np.__version__} hnswlib {hnswlib.__version__ if hasattr(hnswlib,'__version__') else '0.8.0'}")
print("layout,selectivity,method,recall@10,avg_results,median_ms")
for layout in ["random", "topical"]:
    for sel in [0.5, 0.1, 0.01, 0.001]:
        n_t = int(round(1 / sel)) if layout == "random" else None
        if layout == "random":
            tenant = rng.integers(0, n_t, N)
        else:  # whole clusters assigned to tenants
            ntn = int(round(1 / sel))
            tenant = (rng.permutation(CLUSTERS) % ntn)[cl]
        t_id = 0
        mask = tenant == t_id
        ids_t = np.flatnonzero(mask)
        qi = rng.choice(ids_t, NQ, replace=len(ids_t) < NQ)
        qs = X[qi] + 0.3 * rng.normal(size=(NQ, D)).astype(np.float32) / np.sqrt(D)
        qs = (qs / np.linalg.norm(qs, axis=1, keepdims=True)).astype(np.float32)  # keep float32: a float64 query makes numpy upcast the whole matrix
        truth = [exact(mask, q) for q in qs]
        actual = mask.mean()

        def post(kk):
            def f(q):
                lab, _ = index.knn_query(q, k=kk)
                return [l for l in lab[0] if mask[l]][:K]
            return f
        def ingraph(q):
            lab, _ = index.knn_query(q, k=K, filter=lambda i: bool(mask[i]))
            return list(lab[0])
        def pre(q):
            return list(exact(mask, q))
        Xt = X[ids_t]  # tenant partition built once, outside the timed region
        def pre_part(q):
            s = Xt @ q
            top = np.argpartition(-s, K)[:K]
            return list(ids_t[top[np.argsort(-s[top])]])
        for name, fn in [("post k=10", post(K)), ("post k=100", post(100)), ("post k=1000", post(1000)),
                         ("in-graph filter", ingraph), ("pre-filter exact (mask+gather)", pre), ("pre-filter exact (partitioned)", pre_part)]:
            if name == "post k=1000":
                index.set_ef(1000)
            try:
                res, ms = timed(fn, qs)
            except RuntimeError as e:
                print(f"{layout},{actual:.4f},{name},ERROR {str(e)[:60]},,"); continue
            finally:
                index.set_ef(100)
            print(f"{layout},{actual:.4f},{name},{recall(res, truth):.3f},{fill(res):.2f},{ms:.2f}")