"""Filtered search with the HNSW index genuinely forced (no B-tree on the filter column,
enable_seqscan and enable_bitmapscan off), for each iterative_scan mode. Replaces the forced runs in
filtering.json (the first version left enable_seqscan on, so the planner picked seq scan + sort)."""
import json, time
import numpy as np
from common import *

N = 100_000
q = queries()
base = corpus(N)
rng = np.random.default_rng(1)          # same tenant assignment as e6_filter.py
tenant = rng.integers(0, 1000, N)
path = f"{OUT}/filtering.json"
out = json.load(open(path))
out["runs"] = [r for r in out["runs"] if not r.get("forced_hnsw")]

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("VACUUM ANALYZE items_f")
    c.execute("CREATE INDEX items_f_hnsw ON items_f USING hnsw (embedding vector_cosine_ops)")
    c.execute("SET hnsw.ef_search = 40")
    c.execute("SET enable_seqscan = off")
    c.execute("SET enable_bitmapscan = off")
    for sel in (500, 100, 10, 1):
        keep = np.where(tenant < sel)[0]
        gt = keep[ground_truth(base[keep], q, k=10)]
        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] if isinstance(r[0], str) else r[0].decode() for r in c.execute("EXPLAIN (COSTS OFF) " + SQL, (q[0],)).fetchall()]
            assert any("items_f_hnsw" in l for l in plan), plan
            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"], flush=True)
    c.execute("DROP TABLE items_f")
out["forced_note"] = "forced runs: no B-tree on tenant_id, enable_seqscan=off, enable_bitmapscan=off (e6b_forced.py)"
save("filtering", out)