import torch
import pandas as pd
from sentence_transformers import SentenceTransformer, util
import gradio as gr

model = SentenceTransformer("distilbert-base-multilingual-cased")
modela = SentenceTransformer("paraphrase-multilingual-MiniLM-L12-v2")
df = pd.read_csv("cleaned1.csv")
df2 = pd.read_csv("cleaned2.csv")
df3 = pd.read_csv("cleaned3.csv")

embeddings = torch.load("embeddings1_1.pt")
embeddings2 = torch.load("embeddings2_1.pt")
embeddings3 = torch.load("embeddings3_1.pt")

embeddingsa = torch.load("embeddings1.pt")
embeddingsa2 = torch.load("embeddings2.pt")
embeddingsa3 = torch.load("embeddings3.pt")

# Pre-extract DataFrame columns to avoid repeated iloc calls
df_questions = df["question"].values
df_links = df["link"].values
df2_questions = df2["question"].values
df2_links = df2["link"].values
df3_questions = df3["question"].values
df3_links = df3["url"].values
def predict(text):
    if not text or text.strip() == "":
        return "No query provided"
    
    query_embedding = model.encode(text, convert_to_tensor=True)
    query_embeddinga = modela.encode(text, convert_to_tensor=True)
    all_sim_scores1 = []
    all_sim_scores2 = []
    all_sim_scores3 = []
    # Compute similarity scores
    sim_scores1 = util.pytorch_cos_sim(query_embedding, embeddings)[0]
    sim_scores2 = util.pytorch_cos_sim(query_embedding, embeddings2)[0]
    sim_scores3 = util.pytorch_cos_sim(query_embedding, embeddings3)[0]
    
    all_sim_scores1.append(sim_scores1)
    all_sim_scores2.append(sim_scores2)
    all_sim_scores3.append(sim_scores3)
    
    sim_scores1a = util.pytorch_cos_sim(query_embeddinga, embeddingsa)[0]
    sim_scores2a = util.pytorch_cos_sim(query_embeddinga, embeddingsa2)[0]
    sim_scores3a = util.pytorch_cos_sim(query_embeddinga, embeddingsa3)[0]
    
    all_sim_scores1.append(sim_scores1a)
    all_sim_scores2.append(sim_scores2a)
    all_sim_scores3.append(sim_scores3a)
    
    sim_scores1 = torch.stack(all_sim_scores1).mean(dim=0)
    sim_scores2 = torch.stack(all_sim_scores2).mean(dim=0)
    sim_scores3 = torch.stack(all_sim_scores3).mean(dim=0)
    # Get top 3 values and indices in one call
    top3_scores1, top3_idx1 = sim_scores1.topk(3)
    top3_scores2, top3_idx2 = sim_scores2.topk(3)
    top3_scores3, top3_idx3 = sim_scores3.topk(3)
    # Convert to CPU once
    top3_idx1_cpu = top3_idx1.cpu().numpy()
    top3_idx2_cpu = top3_idx2.cpu().numpy()
    top3_idx3_cpu = top3_idx3.cpu().numpy()
    
    top3_scores1_cpu = top3_scores1.cpu().numpy()
    top3_scores2_cpu = top3_scores2.cpu().numpy()
    top3_scores3_cpu = top3_scores3.cpu().numpy()
    # Prepare results using pre-extracted arrays
    results = {
        
        "top2": [
            {
                "question": df2_questions[idx],
                "link": df2_links[idx], 
                "score": float(score)
            }
            for idx, score in zip(top3_idx2_cpu, top3_scores2_cpu)
        ],
        "top3": [
            {
                "question": df3_questions[idx],
                "link": df3_links[idx], 
                "score": float(score)
            }
            for idx, score in zip(top3_idx3_cpu, top3_scores3_cpu)
        ],
        "top1": [
            {
                "question": df_questions[idx],
                "link": df_links[idx],
                "score": float(score)
            }
            for idx, score in zip(top3_idx1_cpu, top3_scores1_cpu)
        ],
    }
    
    return results

# Match the EXACT structure of your working translation app
title = "Search CSV"
iface = gr.Interface(
    fn=predict,  # Changed from search_fatwa to predict 
    inputs=[gr.Textbox(label="text", lines=3)],
    outputs='json',
    title=title,
)
iface.launch()