import gradio as gr
import pandas as pd
import numpy as np
from sentence_transformers import SentenceTransformer
from datasets import load_dataset
from sklearn.metrics.pairwise import cosine_similarity
from PIL import Image

# 1. טעינת המשאבים (Startup)
print("⏳ Loading Model and Data...")
model = SentenceTransformer('clip-ViT-B-32')

# טעינת המטא-דאטה מהקובץ ששמרנו
df = pd.read_parquet("ven_inventory.parquet")
inventory_embeddings = np.stack(df['embedding'].values)

# טעינת התמונות מהדאטה-סט
dataset = load_dataset("detection-datasets/fashionpedia", split='train')
subset = dataset.select(range(5050))

# פונקציית ההמלצה המעודכנת - מקבלת את כל הקלטים
def recommend(text_query, image_query, input_mode):
    if input_mode == "Text":
        if not text_query: return None
        query_emb = model.encode([text_query])
    else:
        if image_query is None: return None
        # המרת התמונה ל-PIL אם היא מגיעה כמערך נומפי
        img = Image.fromarray(image_query).convert("RGB")
        query_emb = model.encode([img])
    
    # נורמליזציה
    query_emb = query_emb / np.linalg.norm(query_emb)
    
    # חישוב דמיון קוסינוס
    scores = cosine_similarity(query_emb, inventory_embeddings)[0]
    top_indices = np.argsort(scores)[::-1][:3]
    
    results = []
    for idx in top_indices:
        actual_idx = int(idx)
        results.append((
            subset[actual_idx]['image'], 
            f"Match Score: {scores[actual_idx]:.2%} | Cluster: {df.iloc[actual_idx]['cluster']}"
        ))
    return results

# 2. בניית ממשק המשתמש (UI)
with gr.Blocks(theme=gr.themes.Soft()) as demo:
    gr.Markdown("# 🌿 Ven Community - Fashion Recommender")
    gr.Markdown("Search Ven's inventory by text or image.")
    
    with gr.Row():
        with gr.Column():
            input_mode = gr.Radio(["Text", "Image"], label="Input Type", value="Text")
            
            # יצירת שני רכיבי הקלט
            text_input = gr.Textbox(label="Description", placeholder="e.g., White sneakers", visible=True)
            image_input = gr.Image(label="Upload Image", visible=False)
            
            search_btn = gr.Button("Find Similar Items", variant="primary")
        
        with gr.Column():
            output_gallery = gr.Gallery(label="Results", columns=3)

    # עדכון נראות הרכיבים לפי הבחירה ב-Radio
    def update_visibility(mode):
        if mode == "Text":
            return gr.update(visible=True), gr.update(visible=False)
        else:
            return gr.update(visible=False), gr.update(visible=True)

    input_mode.change(update_visibility, inputs=input_mode, outputs=[text_input, image_input])
    
    # שליחת כל הרכיבים לפונקציה - זה התיקון הקריטי
    search_btn.click(
        fn=recommend, 
        inputs=[text_input, image_input, input_mode], 
        outputs=output_gallery
    )

demo.launch()