import streamlit as st
from datasets import load_dataset
from PIL import Image
import torch
import numpy as np
from transformers import CLIPProcessor, CLIPModel
from sklearn.metrics.pairwise import cosine_similarity

# =========================
# 📌 Load Model (CLIP)
# =========================
model = CLIPModel.from_pretrained("openai/clip-vit-base-patch32")
processor = CLIPProcessor.from_pretrained("openai/clip-vit-base-patch32")

# =========================
# 📌 Load Dataset (Hugging Face)
# =========================
dataset = load_dataset("fatma812/egypt-artifacts-images")

images = [item["image"] for item in dataset["train"]]

# =========================
# 🧠 Encode Dataset Images
# =========================
@st.cache_resource
def get_image_embeddings():

    embeddings = []

    for img in images:

        inputs = processor(images=img, return_tensors="pt")

        with torch.no_grad():
            emb = model.get_image_features(**inputs)

        embeddings.append(emb[0].numpy())

    return np.array(embeddings)

image_embeddings = get_image_embeddings()

# =========================
# 🎯 UI
# =========================
st.title("🏺 Egyptian Artifacts CLIP Recommender")

st.write("Upload an artifact image and find similar artifacts")

uploaded_file = st.file_uploader("Upload Image", type=["jpg", "png", "jpeg"])

# =========================
# 🔍 Recommendation
# =========================
if uploaded_file:

    query_image = Image.open(uploaded_file).convert("RGB")
    st.image(query_image, caption="Uploaded Image")

    # encode query image
    inputs = processor(images=query_image, return_tensors="pt")

    with torch.no_grad():
        query_emb = model.get_image_features(**inputs).numpy()

    # similarity
    scores = cosine_similarity(query_emb, image_embeddings)[0]

    top_idx = scores.argsort()[::-1][:5]

    st.subheader("🔎 Similar Artifacts")

    for i in top_idx:

        st.image(images[i])
        st.write(f"Artifact #{i}")
        st.write("---")