import datasets
import numpy as np
import scipy.spatial
import scipy.special
import spaces
from sentence_transformers import CrossEncoder, SentenceTransformer

from table import BASE_REPO_ID

ds = datasets.load_dataset(BASE_REPO_ID, split="train")
ds = ds.rename_column("submission_number", "paper_id")
ds.add_faiss_index(column="embedding")

model = SentenceTransformer("all-MiniLM-L6-v2")
reranker = CrossEncoder("cross-encoder/ms-marco-MiniLM-L-6-v2")


@spaces.GPU(duration=5)
def semantic_search(
    query: str, candidate_pool_size: int = 300, score_threshold: float = 0.5
) -> tuple[list[int], list[float]]:
    query_vec = model.encode(query)
    _, retrieved_data = ds.get_nearest_examples("embedding", query_vec, k=candidate_pool_size)

    rerank_inputs = [
        [query, f"{title}\n{abstract}"]
        for title, abstract in zip(retrieved_data["title"], retrieved_data["abstract"], strict=True)
    ]
    rerank_scores = reranker.predict(rerank_inputs)
    sorted_indices = np.argsort(rerank_scores)[::-1]

    paper_ids = []
    scores = []
    for i in sorted_indices:
        score = float(scipy.special.expit(rerank_scores[i]))
        if score < score_threshold:
            break
        paper_ids.append(retrieved_data["paper_id"][i])
        scores.append(score)

    return paper_ids, scores