"""Measure speed and retrieval agreement for BGE sizes on the real question bank."""
import json, sys, time
import numpy as np
from fastembed import TextEmbedding

QUESTIONS = json.load(open("/bench/questions.json"))[:300]
TEXTS = [(q["question_text"] + " " + (q["options"] or ""))[:4000] for q in QUESTIONS]

# Concept queries whose wording deliberately differs from the stems.
QUERIES = [
    "seizure caused by high temperature in a toddler",
    "yellow discolouration in a newborn baby",
    "wheezing infant with runny nose in winter",
    "prolonged fever with red eyes and peeling hands",
    "low iron causing tiredness in a young child",
]

def run(model_name):
    t0 = time.time()
    model = TextEmbedding(model_name=model_name)
    load = time.time() - t0

    t0 = time.time()
    corpus = np.array(list(model.embed(TEXTS)), dtype=np.float32)
    corpus_s = time.time() - t0

    t0 = time.time()
    qvecs = np.array(list(model.embed(QUERIES)), dtype=np.float32)
    query_ms = (time.time() - t0) / len(QUERIES) * 1000

    corpus /= np.linalg.norm(corpus, axis=1, keepdims=True)
    qvecs /= np.linalg.norm(qvecs, axis=1, keepdims=True)
    sims = qvecs @ corpus.T
    top = {}
    for i, q in enumerate(QUERIES):
        idx = np.argsort(-sims[i])[:5]
        top[q] = [(int(QUESTIONS[j]["id"]), round(float(sims[i][j]), 3)) for j in idx]
    return {
        "model": model_name, "dim": corpus.shape[1], "load_s": round(load, 1),
        "corpus_s": round(corpus_s, 1),
        "per_doc_ms": round(corpus_s / len(TEXTS) * 1000, 1),
        "per_query_ms": round(query_ms, 1),
        "top": top,
    }

if __name__ == "__main__":
    results = [run(name) for name in sys.argv[1:]]
    json.dump(results, open("/bench/results.json", "w"), indent=1)
    for r in results:
        print(f"\n{r['model']}  dim={r['dim']}  load={r['load_s']}s  "
              f"{r['per_doc_ms']}ms/doc  {r['per_query_ms']}ms/query  "
              f"(300 docs in {r['corpus_s']}s)")
        for q, hits in r["top"].items():
            print(f"   {q[:46]:48s} top={hits[0]}")