import numpy as np

def cosine(a, b): return float(np.dot(a, b))

def mmr(query_vec, cand_ids, cand_vecs, k=8, lamb=0.7):
    """cand_ids: [int], cand_vecs: np.ndarray float32 [N,D] (unit vectors) aligned with cand_ids"""
    selected, selected_idx = [], []
    remaining = list(range(len(cand_ids)))

    # seed with the most relevant
    best0 = max(remaining, key=lambda i: cosine(query_vec, cand_vecs[i]))
    selected.append(cand_ids[best0]); selected_idx.append(best0); remaining.remove(best0)

    while remaining and len(selected) < k:
        def mmr_score(i):
            rel = cosine(query_vec, cand_vecs[i])
            red = max(cosine(cand_vecs[i], cand_vecs[j]) for j in selected_idx) if selected_idx else 0.0
            return lamb * rel - (1.0 - lamb) * red
        nxt = max(remaining, key=mmr_score)
        selected.append(cand_ids[nxt]); selected_idx.append(nxt); remaining.remove(nxt)
    return selected

def embed_unit_np(st_model, texts: list[str]) -> np.ndarray:
    V = st_model.encode(texts, normalize_embeddings=True, convert_to_numpy=True, batch_size=32)
    V = V.astype("float32", copy=False)
    return V