from rag.constants import BATCH
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=BATCH)
    V = V.astype("float32", copy=False)
    return V

def mmr2(qvec: np.ndarray, ids, vecs: np.ndarray, k=8, lamb=0.7):
    sel, idxs = [], []
    rest = list(range(len(ids)))
    best0 = max(rest, key=lambda i: float(qvec @ vecs[i]))
    sel.append(ids[best0]); idxs.append(best0); rest.remove(best0)
    while rest and len(sel) < k:
        def score(i):
            rel = float(qvec @ vecs[i])
            red = max(float(vecs[i] @ vecs[j]) for j in idxs)
            return lamb*rel - (1-lamb)*red
        nxt = max(rest, key=score)
        sel.append(ids[nxt]); idxs.append(nxt); rest.remove(nxt)
    return sel