"""Hybrid search dataset: 100K DBpedia abstracts embedded locally with BAAI/bge-small-en-v1.5 (384-d),
loaded into Postgres with a generated tsvector column, GIN + HNSW indexes, and a table of query embeddings."""
import json, sys, time
import numpy as np
import psycopg
from pgvector.psycopg import register_vector
from common import NP, save, bge_small

N = 100_000
QUERIES = sys.argv[1:] or []
rows = []
with open(f"{NP}/corpus_meta.jsonl") as fh:
    for i, line in enumerate(fh):
        if i >= N:
            break
        d = json.loads(line)
        rows.append((d["title"], d["text"]))

model = bge_small()
t = time.time()
emb_path = f"{NP}/hybrid_bge384.npy"
try:
    E = np.load(emb_path)
except FileNotFoundError:
    texts = [f"{a}. {b}" for a, b in rows]
    order = sorted(range(N), key=lambda i: len(texts[i]))   # length-sorted batches: far less padding
    E = np.zeros((N, 384), dtype=np.float32)
    for start in range(0, N, 2048):
        idx = order[start:start + 2048]
        E[idx] = np.array(list(model.embed([texts[i] for i in idx], batch_size=32)), dtype=np.float32)
        print(f"  {start + len(idx):,}/{N:,} {time.time() - t:.0f}s", flush=True)
    np.save(emb_path, E)
embed_s = time.time() - t
print("embedded", E.shape, round(embed_s), "s", flush=True)

conn = psycopg.connect(host="/var/run/postgresql", port=5433, user="postgres", dbname="postgres", autocommit=True)
conn.execute("DROP DATABASE IF EXISTS hybrid")
conn.execute("CREATE DATABASE hybrid ENCODING 'UTF8' TEMPLATE template0 LOCALE 'C.UTF-8'")
conn.close()
conn = psycopg.connect(host="/var/run/postgresql", port=5433, user="postgres", dbname="hybrid", autocommit=True)
conn.execute("CREATE EXTENSION vector")
register_vector(conn)
conn.execute("""CREATE TABLE articles (
  id bigint PRIMARY KEY,
  title text NOT NULL,
  body text NOT NULL,
  embedding vector(384) NOT NULL,
  tsv tsvector GENERATED ALWAYS AS (to_tsvector('english', title || ' ' || body)) STORED
)""")
with conn.cursor().copy("COPY articles (id, title, body, embedding) FROM STDIN WITH (FORMAT BINARY)") as cp:
    cp.set_types(["int8", "text", "text", "vector"])
    for i, ((a, b), e) in enumerate(zip(rows, E)):
        cp.write_row([i + 1, a, b, e])
conn.execute("VACUUM ANALYZE articles")
out = {"n": N, "model": "BAAI/bge-small-en-v1.5 (fp32 ONNX, fastembed)", "dims": 384, "embed_s_4cpu": round(embed_s, 1)}
t = time.time(); conn.execute("CREATE INDEX articles_tsv_idx ON articles USING gin (tsv)"); out["gin_build_s"] = round(time.time() - t, 1)
t = time.time(); conn.execute("CREATE INDEX articles_embedding_idx ON articles USING hnsw (embedding vector_cosine_ops)"); out["hnsw_build_s"] = round(time.time() - t, 1)
for rel in ("articles", "articles_tsv_idx", "articles_embedding_idx"):
    out[rel + "_bytes"] = conn.execute("SELECT pg_total_relation_size(%s::regclass)" if rel == "articles" else "SELECT pg_relation_size(%s::regclass)", (rel,)).fetchone()[0]

# query embeddings stored in the database so transcripts are self-contained
conn.execute("CREATE TABLE queries (q text PRIMARY KEY, embedding vector(384))")
qs = QUERIES
for q, e in zip(qs, model.query_embed(qs)):
    conn.execute("INSERT INTO queries VALUES (%s, %s)", (q, np.asarray(e, dtype=np.float32)))
conn.execute("ANALYZE queries")
save("hybrid", out)
print(out)