"""Scaling: HNSW (m=16, ef_construction=64) and exact search at 100K / 250K / 500K rows."""
import numpy as np, time
from common import *

q = queries()
SQL = "SELECT id FROM items_s ORDER BY embedding <=> %s LIMIT 10"
out = {"dim": 1536, "k": 10, "sizes": []}
with connect() as c:
    for n in (100_000, 250_000, 500_000):
        gt = np.load(f"{NP}/gt_{n}.npy")
        r = {"n": n, "load_s": round(load_table(c, "items_s", corpus(n)), 1)}
        r["table_bytes"] = table_size(c, "items_s")
        r["exact_parallel"] = run_queries(c, SQL, q[:100], gt[:100], warmup=5)
        t = time.perf_counter()
        c.execute("CREATE INDEX items_s_hnsw ON items_s USING hnsw (embedding vector_cosine_ops)")
        r["hnsw_build_s"] = round(time.perf_counter() - t, 1)
        r["hnsw_bytes"] = rel_size(c, "items_s_hnsw")
        r["sweep"] = []
        for ef in (40, 100, 200):
            c.execute(f"SET hnsw.ef_search = {ef}")
            s = run_queries(c, SQL, q, gt)
            s["ef_search"] = ef
            r["sweep"].append(s)
        c.execute("RESET hnsw.ef_search")
        print(r, flush=True)
        out["sizes"].append(r)
    c.execute("DROP TABLE items_s")
save("scale", out)