"""IVFFlat: build time, size, recall@10 and latency vs lists / probes."""
import numpy as np, time
from common import *

q = queries()
PROBES = [1, 2, 4, 8, 16, 32, 64]
out = {"dim": 1536, "k": 10, "runs": []}
with connect() as c:
    for n, lists_opts in ((100_000, [100, 316, 1000]), (500_000, [500, 707, 2000])):
        gt = np.load(f"{NP}/gt_{n}.npy")
        load_table(c, "items_i", corpus(n))
        SQL = "SELECT id FROM items_i ORDER BY embedding <=> %s LIMIT 10"
        for lists in lists_opts:
            c.execute("DROP INDEX IF EXISTS items_i_ivf")
            t = time.perf_counter()
            c.execute(f"CREATE INDEX items_i_ivf ON items_i USING ivfflat (embedding vector_cosine_ops) WITH (lists = {lists})")
            r = {"n": n, "lists": lists, "build_s": round(time.perf_counter() - t, 1),
                 "index_bytes": rel_size(c, "items_i_ivf"), "sweep": []}
            for p in PROBES:
                if p > lists:
                    continue
                c.execute(f"SET ivfflat.probes = {p}")
                s = run_queries(c, SQL, q, gt)
                s["probes"] = p
                r["sweep"].append(s)
                print(n, lists, p, s, flush=True)
            c.execute("RESET ivfflat.probes")
            out["runs"].append(r)
    c.execute("DROP TABLE items_i")
save("ivfflat", out)