import pickle
from pathlib import Path

import gradio as gr
from sentence_transformers import CrossEncoder, SentenceTransformer, util

DATA_DIR = Path(__file__).parent
RESULT_COUNT = 5


def load_pickle(filename):
    """Load one of the trusted data files stored with this Space."""
    # Pickle files can execute code while loading. Only load the two files
    # committed to this repository. Never pass user uploads to this function.
    with (DATA_DIR / filename).open("rb") as file:
        return pickle.load(file)


# Load each model once when the Space starts. The corrected reranker ID matches
# the official Sentence Transformers model repository.
bi_encoder = SentenceTransformer("multi-qa-MiniLM-L6-cos-v1")
cross_encoder = CrossEncoder("cross-encoder/ms-marco-MiniLM-L6-v2")
corpus_embeddings = load_pickle("corpus_embeddings_cpu.pkl")
corpus = load_pickle("corpus.pkl")


def search(query, top_k=100):
    """Retrieve relevant passages, rerank them, and return five results."""
    query = (query or "").strip()
    if not query:
        return ("Enter a search query.", "", "", "", "")

    # Semantic search quickly finds a small set of candidate passages.
    # Encode the query using the bi-encoder and find potentially relevant passages
    question_embedding = bi_encoder.encode(query, convert_to_tensor=True)
    hits = util.semantic_search(question_embedding, corpus_embeddings, top_k=top_k)
    hits = hits[0]  # Get the hits for the first query

    # The slower cross-encoder improves result quality by reranking candidates.
    # Now, score all retrieved passages with the cross_encoder
    cross_inp = [[query, corpus[hit['corpus_id']]] for hit in hits]
    cross_scores = cross_encoder.predict(cross_inp)

    # Sort results by the cross-encoder scores
    for hit, score in zip(hits, cross_scores):
        hit["cross_score"] = float(score)

    hits = sorted(hits, key=lambda hit: hit["cross_score"], reverse=True)
    answers = [corpus[hit["corpus_id"]] for hit in hits[:RESULT_COUNT]]

    # Gradio expects one value for each of the five output boxes.
    answers.extend([""] * (RESULT_COUNT - len(answers)))
    return tuple(answers)


examples = [
    "Who is Steve Jobs?",
    "What is Coldplay?",
    "What is a Turing test?",
    "What is the most interesting thing about our universe?",
    "What are the most beautiful places on Earth?",
]

description = (
    "This semantic search engine uses Sentence Transformers to retrieve and "
    "rerank Wikipedia passages. It returns the five most relevant results."
)

search_input = gr.Textbox(lines=1, label="Search query")
outputs = [
    gr.Textbox(type="text", label=f"Search result {index}")
    for index in range(1, RESULT_COUNT + 1)
]

interface = gr.Interface(
    fn=search,
    inputs=search_input,
    outputs=outputs,
    examples=examples,
    article=description,
    title="Neural Search Engine",
)

if __name__ == "__main__":
    interface.launch()