"""FAISS HNSW on the same 100K vectors and queries, for a head-to-head with pgvector HNSW."""
import numpy as np, time, faiss
from common import *

N = 100_000
q = queries()
gt = np.load(f"{NP}/gt_{N}.npy")
base = corpus(N)
out = {"n": N, "dim": 1536, "k": 10, "faiss": faiss.__version__}

faiss.omp_set_num_threads(4)
idx = faiss.IndexHNSWFlat(1536, 16, faiss.METRIC_INNER_PRODUCT)
idx.hnsw.efConstruction = 64
t = time.perf_counter(); idx.add(base); out["build_s_4threads"] = round(time.perf_counter() - t, 1)
out["index_bytes"] = int(len(faiss.serialize_index(idx)))

faiss.omp_set_num_threads(1)
out["sweep"] = []
for ef in (10, 20, 40, 64, 100, 200, 400):
    idx.hnsw.efSearch = ef
    lat, rec = [], []
    for i in range(len(q)):
        t = time.perf_counter()
        _, I = idx.search(q[i:i + 1], 10)
        lat.append((time.perf_counter() - t) * 1000)
        rec.append(recall(list(I[0]), gt[i]))
    lat.sort()
    s = {"ef_search": ef, "recall": round(float(np.mean(rec)), 4), "p50_ms": round(lat[500], 3),
         "p95_ms": round(lat[950], 3), "p99_ms": round(lat[990], 3)}
    out["sweep"].append(s); print(s, flush=True)

flat = faiss.IndexFlatIP(1536); flat.add(base)
lat = []
for i in range(200):
    t = time.perf_counter(); flat.search(q[i:i + 1], 10); lat.append((time.perf_counter() - t) * 1000)
lat.sort(); out["flat_p50_ms_1thread"] = round(lat[100], 3)
save("faiss", out)