"""Top-up for e1: the m=16/ef_construction=64 and m=8 runs whose log lines were truncated,
plus the serial exact-search baseline and plan. Merges into hnsw_100k.json."""
import json, time
import numpy as np
from common import *

N = 100_000
EF = [10, 20, 40, 64, 100, 200, 400]
SQL = "SELECT id FROM items ORDER BY embedding <=> %s LIMIT 10"
q = queries()
gt = np.load(f"{NP}/gt_{N}.npy")
path = f"{OUT}/hnsw_100k.json"
out = json.load(open(path))
with connect() as c:
    n = c.execute("SELECT count(*) FROM items").fetchone()[0]
    assert n == N, n
    plan = c.execute("EXPLAIN (COSTS OFF) " + SQL, (q[0],)).fetchall()
    out["exact_plan"] = [r[0] if isinstance(r[0], str) else r[0].decode() 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")
    for m, efc in ((16, 64), (8, 64)):
        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"] = [x for x in out["builds"] if not (x["m"] == m and x["ef_construction"] == efc)] + [b]
        print(m, efc, b["build_s"], b["index_bytes"], flush=True)
    c.execute("DROP INDEX IF EXISTS items_hnsw")
out["builds"].sort(key=lambda b: (b["m"], b["ef_construction"]))
save("hnsw_100k", out)