import gradio as gr
from datasets import load_dataset
from sentence_transformers import SentenceTransformer
import spaces  # <-- NEW: Import the ZeroGPU library

# 1. Load a fast, lightweight AI model
model = SentenceTransformer('clip-ViT-B-32')

# 2. Load the first 1000 items from your dataset
print("Downloading dataset...")
dataset = load_dataset("JCatesWellcome/sir-henry-wellcome-collection", split="train[:1000]")

# 3. Generate embeddings for the images (Runs on CPU during boot)
print("Analyzing images... (This takes about 2-3 minutes)")
def get_embeddings(batch):
    return {"embeddings": model.encode(batch["image"])}

dataset = dataset.map(get_embeddings, batched=True, batch_size=16)

# 4. Build the searchable index
dataset.add_faiss_index(column="embeddings")
print("App is ready!")

# 5. Define the search function
@spaces.GPU  # <-- NEW: Tell Hugging Face to use the ZeroGPU for this function
def search(query):
    # Convert text to vector
    query_embedding = model.encode(query)
    
    # Find top 6 matches
    scores, samples = dataset.get_nearest_examples("embeddings", query_embedding, k=6)
    
    # Format the output for the Gradio Gallery
    results = []
    for i in range(len(samples["image"])):
        img = samples["image"][i]
        title = samples["title"][i]
        obj_id = samples["object_id"][i]
        caption = f"{title} (ID: {obj_id})"
        results.append((img, caption))
        
    return results

# 6. Build the Web Interface
with gr.Blocks(theme=gr.themes.Soft()) as demo:
    gr.Markdown("# 🏛️ Wellcome Collection Semantic Search")
    gr.Markdown("Search the first 1,000 items of the Sir Henry Wellcome collection using AI. Try searching for concepts, colors, or materials (e.g., 'creepy doll', 'wooden box', 'glass bottle').")
    
    with gr.Row():
        search_box = gr.Textbox(label="What are you looking for?", placeholder="Type here...", scale=4)
        search_button = gr.Button("Search", variant="primary", scale=1)
        
    gallery = gr.Gallery(label="Results", columns=3, height="auto")
    
    # Trigger search on button click or hitting Enter
    search_button.click(fn=search, inputs=search_box, outputs=gallery)
    search_box.submit(fn=search, inputs=search_box, outputs=gallery)

demo.launch()