import pandas as pd
import torch
import torch.nn.functional as F
from transformers import CLIPProcessor, CLIPModel
from PIL import Image
import numpy as np
from sklearn.metrics.pairwise import cosine_similarity
import gradio as gr
import io

# ─── 1. INITIALIZATION ───────────────────────────────────────────
device = "cuda" if torch.cuda.is_available() else "cpu"
MODEL_ID = "openai/clip-vit-base-patch32"

processor = CLIPProcessor.from_pretrained(MODEL_ID)
model = CLIPModel.from_pretrained(MODEL_ID).to(device)
model.eval()

df = pd.read_parquet("food_embeddings.parquet")
embeddings = np.stack(df["embedding"].values)

def idx_to_image(idx):
    return Image.open(io.BytesIO(df.iloc[idx]["image_bytes"]))

print(f"✅ Loaded {len(df)} embeddings from parquet")

# ─── 2. EMBEDDING FUNCTION ───────────────────────────────────────
def get_embedding(user_input, input_type):
    with torch.no_grad():
        if input_type == "image":
            if isinstance(user_input, str):
                image = Image.open(user_input).convert("RGB")
            else:
                image = user_input.convert("RGB")
            inputs = processor(images=image, return_tensors="pt").to(device)
            outputs = model.get_image_features(**inputs)
        else:
            inputs = processor(text=[user_input], return_tensors="pt", padding=True).to(device)
            outputs = model.get_text_features(**inputs)

        if isinstance(outputs, torch.Tensor):
            features = outputs
        else:
            features = outputs.image_embeds if hasattr(outputs, "image_embeds") else outputs.pooler_output

        features = F.normalize(features, p=2, dim=-1)
    return features.cpu().numpy().flatten()

# ─── 3. RECOMMENDATION FUNCTION ─────────────────────────────────
def find_similar_foods(image_input, text_input):
    try:
        if image_input is not None:
            user_vec = get_embedding(image_input, "image")
            mode = "Image"
        elif text_input and text_input.strip():
            user_vec = get_embedding(text_input.strip(), "text")
            mode = "Text"
        else:
            return None, None, None, "⚠️ Please upload an image or enter a text query."

        scores = cosine_similarity(user_vec.reshape(1, -1), embeddings).flatten()
        top3_idx = np.argsort(scores)[::-1][:3]

        results = [idx_to_image(int(idx)) for idx in top3_idx]
        status = f"✅ [{mode} search] Top-3 matches found!"
        return results[0], results[1], results[2], status

    except Exception as e:
        return None, None, None, f"❌ Error: {str(e)}"

# ─── 4. GRADIO UI ────────────────────────────────────────────────
with gr.Blocks(title="Plating Suggestions") as demo:
    gr.Markdown("# 🍽️ Plating Suggestions")
    gr.Markdown("Upload a food photo **or** describe what you're craving — get 3 meal serving suggestions!")

    with gr.Row():
        with gr.Column(scale=1):
            img_input = gr.Image(label="Upload Food Image", type="pil")
            text_input = gr.Textbox(
                label="Or describe a food",
                placeholder="e.g. crispy fried chicken with sauce"
            )
            with gr.Row():
                search_btn = gr.Button("🔍 Find Serving Suggestions", variant="primary")
                clear_btn = gr.Button("✕ Clear")
            status_box = gr.Markdown("")

        with gr.Column(scale=2):
            gr.Markdown("### 🏆 Top-3 Suggestions")
            with gr.Row():
                out1 = gr.Image(label="#1 Best Match")
                out2 = gr.Image(label="#2 Match")
                out3 = gr.Image(label="#3 Match")

    search_btn.click(
        fn=find_similar_foods,
        inputs=[img_input, text_input],
        outputs=[out1, out2, out3, status_box]
    )

    clear_btn.click(
        fn=lambda: (None, "", None, None, None, ""),
        outputs=[img_input, text_input, out1, out2, out3, status_box]
    )

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