"""Embed demo/benchmark queries into hybrid.queries and measure keyword / vector / RRF latency."""
import json, statistics, time
import numpy as np
import psycopg
from pgvector.psycopg import register_vector
from common import OUT, bge_small

QUERIES = [
    "moons of Saturn", "SS-243", "jazz trumpet player from New Orleans", "medieval castle in Lithuania",
    "orchid genus native to South America", "World War II submarine", "Nobel Prize in Chemistry",
    "programming language for the Java virtual machine", "volcano in Ecuador", "marathon runner from Ethiopia",
    "cathedral in Spain", "Canadian foreign minister", "South Korean film director", "Hungarian pianist",
    "parabolic reflector", "railway station in Essex", "football club founded in 1899", "rock band from Liverpool",
    "vaccine for tuberculosis", "chess opening", "Formula One driver", "village in Poland", "Roman emperor",
    "species of moth", "American politician from Ohio", "album released in 1987", "river in Illinois",
    "Russian grand duchess", "national park in Australia", "web service protocol",
]
model = bge_small()
conn = psycopg.connect(host="/var/run/postgresql", port=5433, user="postgres", dbname="hybrid", autocommit=True)
register_vector(conn)
conn.execute("DELETE FROM queries")
for q, e in zip(QUERIES, model.query_embed(QUERIES)):
    conn.execute("INSERT INTO queries VALUES (%s, %s)", (q, np.asarray(e, dtype=np.float32)))
conn.execute("ANALYZE queries")

KW = """SELECT id FROM articles, websearch_to_tsquery('english', %(q)s) query
        WHERE tsv @@ query ORDER BY ts_rank_cd(tsv, query) DESC LIMIT 10"""
VEC = """SELECT id FROM articles ORDER BY embedding <=> (SELECT embedding FROM queries WHERE q = %(q)s) LIMIT 10"""
HYB = """WITH semantic AS (
  SELECT id, row_number() OVER (ORDER BY embedding <=> (SELECT embedding FROM queries WHERE q = %(q)s)) AS rank
  FROM articles ORDER BY embedding <=> (SELECT embedding FROM queries WHERE q = %(q)s) LIMIT 20),
keyword AS (
  SELECT id, row_number() OVER (ORDER BY ts_rank_cd(tsv, query) DESC) AS rank
  FROM articles, websearch_to_tsquery('english', %(q)s) query
  WHERE tsv @@ query ORDER BY ts_rank_cd(tsv, query) DESC LIMIT 20)
SELECT id, COALESCE(1.0 / (60 + s.rank), 0) + COALESCE(1.0 / (60 + k.rank), 0) AS score
FROM semantic s FULL OUTER JOIN keyword k USING (id) ORDER BY score DESC LIMIT 10"""
out = {"queries": QUERIES, "lat": {}}
cur = conn.cursor()
for name, sql in (("keyword", KW), ("vector", VEC), ("hybrid", HYB)):
    for q in QUERIES[:5]:
        cur.execute(sql, {"q": q}).fetchall()
    lat, nres = [], []
    for rep in range(5):
        for q in QUERIES:
            t = time.perf_counter(); rows = cur.execute(sql, {"q": q}).fetchall(); lat.append((time.perf_counter() - t) * 1000)
            nres.append(len(rows))
    lat.sort()
    out["lat"][name] = {"p50_ms": round(lat[len(lat) // 2], 3), "p95_ms": round(lat[int(len(lat) * 0.95)], 3), "avg_results": round(statistics.fmean(nres), 2)}
    print(name, out["lat"][name], flush=True)
# overlap between methods on the demo queries
ov = []
for q in QUERIES:
    k = {r[0] for r in cur.execute(KW, {"q": q}).fetchall()}
    v = {r[0] for r in cur.execute(VEC, {"q": q}).fetchall()}
    ov.append(len(k & v))
out["mean_overlap_kw_vec_top10"] = round(statistics.fmean(ov), 2)
prev = json.load(open(f"{OUT}/hybrid.json"))
prev.update(out)
json.dump(prev, open(f"{OUT}/hybrid.json", "w"), indent=1)
print("saved", out["mean_overlap_kw_vec_top10"])