import gradio as gr
import numpy as np
import pandas as pd
import torch
from datasets import load_dataset
from transformers import CLIPProcessor, CLIPModel

# Load dataset
ds = load_dataset("amaye15/landscapes")
train_ds = ds["train"]

# Load embeddings
df = pd.read_parquet("image_embeddings_clip.parquet")
image_indices = df["image_index"].values
emb_matrix = df.drop(columns=["image_index"]).values.astype(np.float32)

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

def l2_normalize(x):
    return x / np.linalg.norm(x)

@torch.no_grad()
def embed_image(img):
    inputs = processor(images=img, return_tensors="pt")
    inputs = {k: v.to(device) for k, v in inputs.items()}
    feats = model.get_image_features(**inputs)
    feats = feats / feats.norm(dim=-1, keepdim=True)
    return feats.squeeze(0).cpu().numpy()

def recommend(img):
    q_emb = embed_image(img)
    sims = emb_matrix @ l2_normalize(q_emb)
    top = np.argsort(-sims)[1:4]
    results = []
    for i in top:
        results.append(train_ds[int(image_indices[i])]["pixel_values"])
    return results

# Gradio interface
demo = gr.Interface(
    fn=recommend,
    inputs=gr.Image(type="pil", label="Upload a landscape image"),
    outputs=[
        gr.Image(label="Recommendation 1"),
        gr.Image(label="Recommendation 2"),
        gr.Image(label="Recommendation 3"),
    ],
    title="Landscape Image Recommendation System",
    description="Upload a landscape image and receive visually similar recommendations."
)

# App layout with video below
with gr.Blocks() as app:
    demo.render()

    gr.Markdown("---")
    gr.Markdown("## Project Presentation Video")

    gr.HTML("""
    <div style="width:100%;max-width:900px;margin:0 auto;">
      <div style="position:relative;padding-bottom:56.25%;height:0;overflow:hidden;border-radius:12px;">
        <iframe
          src="https://www.youtube.com/embed/nwbPAR7UApw"
          title="Project presentation"
          frameborder="0"
          allow="accelerometer; autoplay; clipboard-write; encrypted-media; gyroscope; picture-in-picture; web-share"
          allowfullscreen
          style="position:absolute;top:0;left:0;width:100%;height:100%;">
        </iframe>
      </div>
    </div>
    """)

app.launch()