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

# Model
model = SentenceTransformer("all-MiniLM-L6-v2")

# Load and process supplier data on startup
def load_supplier_data():
    global supplier_df, supplier_embeddings

    supplier_df = pd.read_excel("SupplierList.xlsx")
    supplier_df["Capability"] = supplier_df["Capability"].fillna("")
    texts = supplier_df["Capability"].tolist()

    # Compute and normalize embeddings
    supplier_embeddings = model.encode(texts, convert_to_tensor=True, normalize_embeddings=True)

# Initial data load
supplier_df = None
supplier_embeddings = None
load_supplier_data()

# View loaded data
def view_supplier_data():
    return supplier_df if supplier_df is not None else pd.DataFrame()

# Match function
def match_suppliers(event_description, top_k=5):
    if supplier_embeddings is None:
        return "Supplier data not loaded."

    event_embedding = model.encode(event_description, convert_to_tensor=True, normalize_embeddings=True).unsqueeze(0)
    scores = torch.nn.functional.cosine_similarity(event_embedding, supplier_embeddings)
    top_indices = torch.topk(scores, k=top_k).indices.tolist()

    results = []
    for idx in top_indices:
        results.append({
            "Supplier Name": supplier_df.loc[idx, "Supplier Name"],
            "Match Score": round(scores[idx].item(), 4),
            "Capabilities": supplier_df.loc[idx, "Capability"]
        })

    return pd.DataFrame(results)

# Gradio Interface
with gr.Blocks() as demo:
    gr.Markdown("## Supplier Matching App")

    with gr.Tab("2. Match Use Case to Suppliers"):
        input_text = gr.Textbox(lines=3, label="Describe Your Event/Use Case")
        top_k_input = gr.Slider(minimum=1, maximum=10, value=5, label="Top K Matches")
        match_button = gr.Button("Match Suppliers")
        match_output = gr.Dataframe()
        match_button.click(match_suppliers, inputs=[input_text, top_k_input], outputs=match_output)

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