import json
import math
import gradio as gr
import torch
from sentence_transformers import SentenceTransformer, CrossEncoder, util

HF_USERNAME = "anuragseven"  # <-- set to your username, or the Space will fail to load

# fine-tuned models (downloaded once from the Hub, then cached on the Space)
bi = SentenceTransformer(f"{HF_USERNAME}/scifact-biencoder")
ce = CrossEncoder(f"{HF_USERNAME}/scifact-crossencoder")

# precomputed artifacts shipped in this repo -> no beir, no corpus download,
# no startup embedding. Cold start is just "load three small files".
with open("corpus.json") as f:
    corpus = json.load(f)          # {doc_id: {"title": ..., "text": ...}}
with open("ids.json") as f:
    ids = json.load(f)             # ordered doc_ids aligned with emb rows
emb = torch.load("doc_emb.pt", map_location="cpu")  # normalized doc embeddings

RERANK_K = 20  # dense candidates fed to the cross-encoder (was 50)


def make_text(d):
    return (corpus[d].get("title", "") + " " + corpus[d].get("text", "")).strip()


def search(query, top_k=5):
    if not query or not query.strip():
        return "Enter a scientific claim or query."
    qe = bi.encode(query, normalize_embeddings=True, convert_to_tensor=True)
    hits = util.semantic_search(qe, emb, top_k=RERANK_K)[0]
    cand = [ids[h["corpus_id"]] for h in hits]
    scores = ce.predict([[query, make_text(c)] for c in cand])
    order = sorted(range(len(cand)), key=lambda j: scores[j], reverse=True)[:int(top_k)]
    out = ""
    for r, j in enumerate(order, 1):
        c = cand[j]
        rel = 1.0 / (1.0 + math.exp(-float(scores[j])))  # sigmoid: logit -> 0..1
        out += f"### {r}. {corpus[c]['title']}\n{corpus[c]['text'][:300]}...\n\n_relevance: {rel:.2f}_\n\n"
    return out


EXAMPLES = [
    ["Statins reduce LDL cholesterol levels.", 5],
    ["Aspirin reduces the risk of colorectal cancer.", 5],
    ["Vitamin D supplementation reduces respiratory infections.", 5],
    ["Smoking increases the risk of cardiovascular disease.", 5],
]


gr.Interface(fn=search,
             inputs=[gr.Textbox(label="Scientific claim / query"), gr.Slider(1, 10, value=5, step=1, label="Results")],
             outputs=gr.Markdown(),
             examples=EXAMPLES,
             cache_examples=False,  # click only fills the box; user presses Submit (clear feedback)
             title="SciFact Neural Search",
             description="Fine-tuned bi-encoder retrieval + cross-encoder re-ranking").launch(server_name="0.0.0.0", server_port=7860)