"""HNSW: build time, size, recall@10 and latency vs m / ef_construction / ef_search (100K x 1536)."""
import numpy as np, time
from common import *

N = 100_000
EF = [10, 20, 40, 64, 100, 200, 400]
BUILDS = [(16, 64), (8, 64), (32, 64), (16, 32), (16, 128), (16, 256)]
SQL = "SELECT id FROM items ORDER BY embedding <=> %s LIMIT 10"

q = queries()
gt = np.load(f"{NP}/gt_{N}.npy")
out = {"n": N, "dim": 1536, "k": 10, "queries": len(q), "metric": "cosine (vector_cosine_ops)"}

with connect() as c:
    out["load_s"] = round(load_table(c, "items", corpus(N)), 1)
    out["table_bytes"] = table_size(c, "items")

    # Exact search baseline (no index): parallel seq scan, and single process.
    c.execute("SET max_parallel_workers_per_gather = 2")
    out["exact_parallel"] = run_queries(c, SQL, q[:200], gt[:200], warmup=10)
    plan = c.execute("EXPLAIN (COSTS OFF) " + SQL.replace("%s", "'[" + ",".join(["0"] * 1535 + ["1"]) + "]'")).fetchall()
    out["exact_plan"] = [r[0] for r in plan]
    c.execute("SET max_parallel_workers_per_gather = 0")
    out["exact_serial"] = run_queries(c, SQL, q[:200], gt[:200], warmup=10)
    c.execute("RESET max_parallel_workers_per_gather")

    out["builds"] = []
    for m, efc in BUILDS:
        c.execute("DROP INDEX IF EXISTS items_hnsw")
        t = time.perf_counter()
        c.execute(f"CREATE INDEX items_hnsw ON items USING hnsw (embedding vector_cosine_ops) WITH (m = {m}, ef_construction = {efc})")
        b = {"m": m, "ef_construction": efc, "build_s": round(time.perf_counter() - t, 1),
             "index_bytes": rel_size(c, "items_hnsw"), "sweep": []}
        for ef in EF:
            c.execute(f"SET hnsw.ef_search = {ef}")
            r = run_queries(c, SQL, q, gt)
            r["ef_search"] = ef
            b["sweep"].append(r)
            print(m, efc, ef, r, flush=True)
        c.execute("RESET hnsw.ef_search")
        out["builds"].append(b)
        print(b["m"], b["ef_construction"], b["build_s"], b["index_bytes"], flush=True)
    c.execute("DROP INDEX IF EXISTS items_hnsw")
save("hnsw_100k", out)