"""halfvec and binary quantization vs full-precision vector (100K x 1536, HNSW m=16 ef_construction=64)."""
import numpy as np, time
from common import *

N = 100_000
q = queries()
gt = np.load(f"{NP}/gt_{N}.npy")
EF = [10, 20, 40, 64, 100, 200, 400]
out = {"n": N, "dim": 1536, "k": 10}

with connect() as c:
    # --- halfvec column -------------------------------------------------------------
    out["halfvec_load_s"] = round(load_table(c, "items_h", corpus(N), coltype="halfvec"), 1)
    out["halfvec_table_bytes"] = table_size(c, "items_h")
    t = time.perf_counter()
    c.execute("CREATE INDEX items_h_hnsw ON items_h USING hnsw (embedding halfvec_cosine_ops)")
    out["halfvec_build_s"] = round(time.perf_counter() - t, 1)
    out["halfvec_index_bytes"] = rel_size(c, "items_h_hnsw")
    SQLH = "SELECT id FROM items_h ORDER BY embedding <=> %s LIMIT 10"
    out["halfvec_sweep"] = []
    for ef in EF:
        c.execute(f"SET hnsw.ef_search = {ef}")
        s = run_queries(c, SQLH, q, gt, wrap=HalfVector)
        s["ef_search"] = ef
        out["halfvec_sweep"].append(s)
        print("halfvec", ef, s, flush=True)
    c.execute("RESET hnsw.ef_search")
    c.execute("DROP TABLE items_h")

    # --- binary quantization on a full-precision table --------------------------------
    load_table(c, "items_b", corpus(N))
    out["vector_table_bytes"] = table_size(c, "items_b")
    t = time.perf_counter()
    c.execute("CREATE INDEX items_b_bq ON items_b USING hnsw ((binary_quantize(embedding)::bit(1536)) bit_hamming_ops)")
    out["bq_build_s"] = round(time.perf_counter() - t, 1)
    out["bq_index_bytes"] = rel_size(c, "items_b_bq")

    # Without re-ranking: top 10 by Hamming distance only.
    SQL0 = "SELECT id FROM items_b ORDER BY binary_quantize(embedding)::bit(1536) <~> binary_quantize(%s) LIMIT 10"
    c.execute("SET hnsw.ef_search = 100")
    out["bq_no_rerank"] = run_queries(c, SQL0, q, gt)
    print("bq no rerank", out["bq_no_rerank"], flush=True)

    # With re-ranking: fetch C candidates by Hamming, re-order by exact cosine.
    out["bq_rerank"] = []
    for cand in (20, 40, 80, 160, 320, 640):
        c.execute(f"SET hnsw.ef_search = {max(cand, 40)}")
        SQLR = f"""SELECT id FROM (
              SELECT id, embedding FROM items_b
              ORDER BY binary_quantize(embedding)::bit(1536) <~> binary_quantize(%s)
              LIMIT {cand}) c
            ORDER BY c.embedding <=> %s LIMIT 10"""
        cur = c.cursor()
        lat, rec = [], []
        for v in q[:50]:
            cur.execute(SQLR, (v, v), prepare=True).fetchall()
        for i, v in enumerate(q):
            t = time.perf_counter()
            ids = [r[0] for r in cur.execute(SQLR, (v, v), prepare=True).fetchall()]
            lat.append((time.perf_counter() - t) * 1000)
            rec.append(recall(ids, gt[i]))
        lat.sort()
        s = {"candidates": cand, "ef_search": max(cand, 40), "recall": round(float(np.mean(rec)), 4),
             "p50_ms": round(lat[500], 3), "p95_ms": round(lat[950], 3), "p99_ms": round(lat[990], 3)}
        out["bq_rerank"].append(s)
        print("bq rerank", s, flush=True)
    c.execute("RESET hnsw.ef_search")

    # halfvec expression index over the full-precision column (functional check + size).
    t = time.perf_counter()
    c.execute("CREATE INDEX items_b_hexpr ON items_b USING hnsw ((embedding::halfvec(1536)) halfvec_cosine_ops)")
    out["halfvec_expr_build_s"] = round(time.perf_counter() - t, 1)
    out["halfvec_expr_index_bytes"] = rel_size(c, "items_b_hexpr")
    c.execute("SET hnsw.ef_search = 40")
    out["halfvec_expr_ef40"] = run_queries(
        c, "SELECT id FROM items_b ORDER BY embedding::halfvec(1536) <=> %s::halfvec(1536) LIMIT 10", q, gt)
    out["halfvec_expr_plan"] = [r[0] for r in c.execute(
        "EXPLAIN (COSTS OFF) SELECT id FROM items_b ORDER BY embedding::halfvec(1536) <=> %s::halfvec(1536) LIMIT 10",
        (q[0],)).fetchall()]
    c.execute("DROP TABLE items_b")
save("quantization", out)