"""Labelled retrieval test: does a question's own explanation retrieve that question?

The explanation restates the case in different words, so it is a genuine
paraphrase query with a known correct answer — no hand labelling needed.
Reports recall@1 / @5 and mean reciprocal rank against a 500-question corpus.
"""
import json, sys, time
import numpy as np
from fastembed import TextEmbedding

ROWS = json.load(open("/bench/labelled.json"))[:150]
DOCS = [(r["question_text"] + " " + (r["options"] or ""))[:4000] for r in ROWS]
# Query = explanation only, truncated, so wording differs from the stem.
QUERIES = [r["explanation"][:600] for r in ROWS]

def run(name):
    model = TextEmbedding(model_name=name)
    t0 = time.time()
    docs = np.array(list(model.embed(DOCS)), dtype=np.float32)
    doc_s = time.time() - t0
    t0 = time.time()
    qs = np.array(list(model.embed(QUERIES)), dtype=np.float32)
    q_ms = (time.time() - t0) / len(QUERIES) * 1000

    docs /= np.linalg.norm(docs, axis=1, keepdims=True)
    qs /= np.linalg.norm(qs, axis=1, keepdims=True)
    sims = qs @ docs.T
    order = np.argsort(-sims, axis=1)
    gold = np.arange(len(ROWS))
    ranks = np.array([np.where(order[i] == gold[i])[0][0] + 1 for i in range(len(ROWS))])
    return {
        "model": name, "dim": docs.shape[1], "n": len(ROWS),
        "recall@1": round(float((ranks == 1).mean()), 3),
        "recall@5": round(float((ranks <= 5).mean()), 3),
        "mrr": round(float((1 / ranks).mean()), 3),
        "ms_per_doc": round(doc_s / len(DOCS) * 1000, 1),
        "ms_per_query": round(q_ms, 1),
    }

if __name__ == "__main__":
    for name in sys.argv[1:]:
        r = run(name)
        print(f"{r['model']:28s} dim={r['dim']:5d} R@1={r['recall@1']:.3f} "
              f"R@5={r['recall@5']:.3f} MRR={r['mrr']:.3f} "
              f"{r['ms_per_doc']}ms/doc {r['ms_per_query']}ms/query", flush=True)