"""pgvector (seq scan exact, HNSW) vs numpy, 100k x 384 clustered normalized vectors.
Needs: docker run -d --shm-size=1g --name pg -e POSTGRES_PASSWORD=pw -p 5433:5432 pgvector/pgvector:pg17"""
import time, sys, numpy as np, psycopg
from pgvector.psycopg import register_vector
N, D, K, NQ = 100_000, 384, 10, 200
rng = np.random.default_rng(1)
def norm(x): return x / np.linalg.norm(x, axis=1, keepdims=True)
C = norm(rng.standard_normal((2000, D), dtype=np.float32))
lab = rng.integers(0, 2000, N)
X = norm(C[lab] + 0.03 * rng.standard_normal((N, D), dtype=np.float32)).astype(np.float32)
Q = norm(X[rng.integers(0, N, NQ)] + 0.05 * rng.standard_normal((NQ, D), dtype=np.float32)).astype(np.float32)
gt = np.argsort(-(Q @ X.T), axis=1)[:, :K]

def p50(f):
    t = []; out = []
    for q in Q:
        s = time.perf_counter(); r = f(q); t.append((time.perf_counter() - s) * 1e3); out.append(r)
    return np.median(t), np.percentile(t, 95), out
def recall(out): return np.mean([len(set(a) & set(b)) / K for a, b in zip(out, gt)])

def np_search(q):
    s = X @ q; i = np.argpartition(-s, K)[:K]; return list(i[np.argsort(-s[i])])
for _ in range(5): np_search(Q[0])
m, p95, o = p50(np_search); print(f"numpy exact                      p50={m:6.2f}ms p95={p95:6.2f}ms recall@10={recall(o):.3f}", flush=True)

con = psycopg.connect("postgresql://postgres:pw@localhost:5433/postgres", autocommit=True)
con.execute("create extension if not exists vector")
register_vector(con)
con.execute("drop table if exists items"); con.execute("create table items (id int primary key, emb vector(384))")
s = time.perf_counter()
with con.cursor().copy("copy items (id, emb) from stdin (format binary)") as cp:
    cp.set_types(["int4", "vector"])
    for i in range(N): cp.write_row((i, X[i]))
print(f"load {N} rows: {time.perf_counter()-s:.1f}s", flush=True)
con.execute("vacuum analyze items")
print("table size:", con.execute("select pg_size_pretty(pg_table_size('items'))").fetchone()[0])

def pg_search(q):
    return [r[0] for r in con.execute("select id from items order by emb <=> %s limit 10", (q,), prepare=True).fetchall()]
def run(label):
    for _ in range(5): pg_search(Q[0])
    m, p95, o = p50(pg_search); print(f"{label:32} p50={m:6.2f}ms p95={p95:6.2f}ms recall@10={recall(o):.3f}", flush=True)

for w in (0, 2):
    con.execute(f"set max_parallel_workers_per_gather = {w}")
    run(f"pgvector seq scan, workers={w}")
con.execute("set max_parallel_workers_per_gather = 2")
con.execute("set maintenance_work_mem = '1GB'")
s = time.perf_counter()
con.execute("create index items_hnsw on items using hnsw (emb vector_cosine_ops) with (m=16, ef_construction=64)")
print(f"HNSW build (m=16, efc=64): {time.perf_counter()-s:.1f}s; index size:", con.execute("select pg_size_pretty(pg_relation_size('items_hnsw'))").fetchone()[0], flush=True)
for ef in (10, 40, 100, 200):
    con.execute(f"set hnsw.ef_search = {ef}")
    run(f"pgvector HNSW ef_search={ef}")