"""Filtered vector search: HNSW with and without iterative index scans, at several selectivities.

Each row gets a tenant_id uniformly in 0..999; a filter `tenant_id < s` keeps s/1000 of rows.
"""
import numpy as np, time
from common import *

N = 100_000
q = queries()
base = corpus(N)
rng = np.random.default_rng(1)
tenant = rng.integers(0, 1000, N)
out = {"n": N, "dim": 1536, "k": 10, "hnsw": "m=16, ef_construction=64", "ef_search": 40, "runs": []}

with connect() as c:
    load_table(c, "items_f", base, extra_cols=", tenant_id int")
    with c.cursor() as cur:
        cur.execute("CREATE TEMP TABLE t (id int, tenant_id int)")
        with cur.copy("COPY t FROM STDIN") as cp:
            for i, tt in enumerate(tenant):
                cp.write_row((i, int(tt)))
        cur.execute("UPDATE items_f SET tenant_id = t.tenant_id FROM t WHERE t.id = items_f.id")
    c.execute("CREATE INDEX ON items_f (tenant_id)")
    c.execute("VACUUM ANALYZE items_f")
    t = time.perf_counter()
    c.execute("CREATE INDEX items_f_hnsw ON items_f USING hnsw (embedding vector_cosine_ops)")
    out["build_s"] = round(time.perf_counter() - t, 1)
    c.execute("SET hnsw.ef_search = 40")

    for sel in (500, 100, 10, 1):  # 50%, 10%, 1%, 0.1%
        keep = np.where(tenant < sel)[0]
        g = ground_truth(base[keep], q, k=10)
        gt = keep[g]  # map back to ids
        SQL = f"SELECT id FROM items_f WHERE tenant_id < {sel} ORDER BY embedding <=> %s LIMIT 10"
        for mode in ("off", "relaxed_order", "strict_order"):
            c.execute(f"SET hnsw.iterative_scan = {mode}")
            plan = [r[0] for r in c.execute("EXPLAIN (COSTS OFF) " + SQL, (q[0],)).fetchall()]
            r = run_queries(c, SQL, q[:300], gt[:300], warmup=20)
            r.update({"selectivity_pct": sel / 10, "iterative_scan": mode, "plan": plan,
                      "rows_matching": int(len(keep))})
            out["runs"].append(r)
            print(sel, mode, r["recall"], r["avg_rows"], r["p50_ms"], plan[0], flush=True)
        # Force the HNSW index (disable the B-tree path) to show what the index alone does.
        c.execute("SET enable_bitmapscan = off"); c.execute("SET enable_indexscan = on")
        c.execute("SET hnsw.iterative_scan = off")
        c.execute("DROP INDEX IF EXISTS items_f_tenant_id_idx")
        for mode in ("off", "relaxed_order"):
            c.execute(f"SET hnsw.iterative_scan = {mode}")
            plan = [r[0] for r in c.execute("EXPLAIN (COSTS OFF) " + SQL, (q[0],)).fetchall()]
            r = run_queries(c, SQL, q[:300], gt[:300], warmup=20)
            r.update({"selectivity_pct": sel / 10, "iterative_scan": mode, "plan": plan, "forced_hnsw": True,
                      "rows_matching": int(len(keep))})
            out["runs"].append(r)
            print("forced", sel, mode, r["recall"], r["avg_rows"], r["p50_ms"], plan[0], flush=True)
        c.execute("RESET enable_bitmapscan")
        c.execute("CREATE INDEX items_f_tenant_id_idx ON items_f (tenant_id)")
        c.execute("ANALYZE items_f")
    c.execute("RESET hnsw.iterative_scan")
    c.execute("DROP TABLE items_f")
save("filtering", out)