import gradio as gr
import numpy as np
import pandas as pd
import torch
from transformers import CLIPProcessor, CLIPModel
from sklearn.metrics.pairwise import cosine_similarity
from datasets import load_dataset

# Load model
device = "cuda" if torch.cuda.is_available() else "cpu"
model_name = "openai/clip-vit-base-patch32"
model = CLIPModel.from_pretrained(model_name).to(device)
processor = CLIPProcessor.from_pretrained(model_name)

# Load embeddings
df = pd.read_parquet("nature_embeddings.parquet")
EMBEDDINGS_MATRIX = np.array(df['embedding'].tolist())
CAPTIONS = df['caption'].tolist()
ORIGINAL_INDICES = df['original_index'].tolist()

# Load dataset
dataset = load_dataset("mertcobanov/nature-dataset")
train_data = dataset['train']

def classify_scene(caption):
    """Classify nature scene based on caption keywords"""
    caption = caption.lower()
    if any(w in caption for w in ['beach', 'ocean', 'sea', 'coast', 'wave']):
        return '🌊 Coastal'
    elif any(w in caption for w in ['waterfall', 'stream', 'river', 'lake']):
        return '💧 Water'
    elif any(w in caption for w in ['mountain', 'peak', 'cliff', 'rock']):
        return '🏔️ Mountain'
    elif any(w in caption for w in ['desert', 'dune', 'sand', 'arid']):
        return '🏜️ Desert'
    elif any(w in caption for w in ['forest', 'tree', 'jungle', 'wood']):
        return '🌲 Forest'
    else:
        return '🌿 Nature'

def get_image_embedding(image):
    inputs = processor(images=image, return_tensors="pt", padding=True).to(device)
    with torch.no_grad():
        outputs = model.vision_model(**inputs)
        features = outputs.pooler_output
        features = torch.nn.functional.normalize(features, p=2, dim=-1)
    return features.cpu().numpy()

def recommend(image):
    query_embedding = get_image_embedding(image)
    similarities = cosine_similarity(query_embedding, EMBEDDINGS_MATRIX)[0]
    top_indices = np.argsort(similarities)[::-1]
    top_indices = [idx for idx in top_indices if similarities[idx] < 0.9999][:3]
    
    results = []
    for idx in top_indices:
        original_idx = ORIGINAL_INDICES[idx]
        img = train_data[int(original_idx)]['image']
        cap = CAPTIONS[idx]
        sim = float(similarities[idx])
        scene = classify_scene(cap)
        results.append((img, f"{scene} | {cap} (similarity: {sim:.4f})"))
    
    return results[0][0], results[0][1], results[1][0], results[1][1], results[2][0], results[2][1]

# Gradio interface
with gr.Blocks(title="Nature Scene Recommender") as demo:
    gr.Markdown("# 🌿 Nature Scene Recommender")
    gr.Markdown("Upload a nature image and get the 3 most similar scenes!")
    
    with gr.Row():
        input_image = gr.Image(type="pil", label="Upload Nature Image")
        btn = gr.Button("Find Similar Scenes 🔍", variant="primary")
    
    gr.Markdown("### Top 3 Similar Scenes")
    with gr.Row():
        img1 = gr.Image(label="Match #1")
        img2 = gr.Image(label="Match #2")
        img3 = gr.Image(label="Match #3")
    with gr.Row():
        cap1 = gr.Textbox(label="Scene #1")
        cap2 = gr.Textbox(label="Scene #2")
        cap3 = gr.Textbox(label="Scene #3")
    
    btn.click(fn=recommend, inputs=input_image,
              outputs=[img1, cap1, img2, cap2, img3, cap3])

demo.launch(server_name="0.0.0.0", server_port=7860)