import gradio as gr
import json
from transformers import AutoProcessor, AutoModelForImageClassification
from PIL import Image
import torch

# Load food DB
with open("merged_foods_with_fruits (1).json", "r", encoding="utf-8") as f:
    data = json.load(f)

def all_foods():
    foods = []
    for category, items in data.items():
        foods.extend(items)
    return foods

def find_food(food_name: str):
    matches = []
    for item in all_foods():
        if food_name in item["نام"]:
            matches.append(item)
    return matches

# Load AI Model
model_name = "eslamxm/vit-base-food101"
processor = AutoProcessor.from_pretrained(model_name)
model = AutoModelForImageClassification.from_pretrained(model_name)

# English→Persian dictionary
english_to_persian = {
    "pizza": "پیتزا",
    "cheeseburger": "همبرگر",
    "burger": "همبرگر",
    "hamburger": "همبرگر",
    "hot dog": "هات داگ",
    "sandwich": "ساندویچ",
    "french fries": "سیب‌زمینی سرخ‌کرده",
    "fried chicken": "مرغ سوخاری",
    "shawarma": "شاورما",
    "kebab": "کباب",
    "rice": "برنج",
    "pasta": "پاستا",
    "sushi": "سوشی",
    "salad": "سالاد",
    "steak": "استیک",
    "chicken": "مرغ",
    "fish": "ماهی",
    "apple": "سیب",
    "banana": "موز",
    "orange": "پرتقال",
}

# Text search
def analyze_text(food):
    results = find_food(food)
    if results:
        return "\n\n".join(
            [f"🍽 {item['نام']} → 🔥 {item['کالری']} کالری | 💪 {item['پروتئین']} | 🍞 {item['کربوهیدرات']} | 🥓 {item['چربی']}" for item in results]
        )
    else:
        return "❌ غذا پیدا نشد"

# Image search
def analyze_image(img):
    image = Image.fromarray(img).convert("RGB")
    inputs = processor(images=image, return_tensors="pt")
    with torch.no_grad():
        outputs = model(**inputs)
        logits = outputs.logits
        predicted_class_id = logits.argmax(-1).item()

    english_label = model.config.id2label[predicted_class_id]
    persian_label = english_to_persian.get(english_label.lower(), english_label)

    results = find_food(persian_label)
    if results:
        return f"🍔 Prediction: {english_label} → {persian_label}\n\n" + "\n\n".join(
            [f"🍽 {item['نام']} → 🔥 {item['کالری']} کالری | 💪 {item['پروتئین']} | 🍞 {item['کربوهیدرات']} | 🥓 {item['چربی']}" for item in results]
        )
    else:
        return f"🍔 Prediction: {english_label} → {persian_label}\n❌ غذا در دیتابیس یافت نشد"

# Gradio UI
with gr.Blocks() as demo:
    gr.Markdown("## 🍽 Persian Food AI (نسخه دمو)")
    
    with gr.Tab("🔎 جستجو با متن"):
        food_input = gr.Textbox(label="نام غذا (مثال: پیتزا)")
        text_output = gr.Textbox(label="نتایج", lines=8)
        food_btn = gr.Button("جستجو")
        food_btn.click(fn=analyze_text, inputs=food_input, outputs=text_output)

    with gr.Tab("📷 تشخیص از عکس"):
        img_input = gr.Image(type="numpy", label="بارگذاری عکس غذا")
        img_output = gr.Textbox(label="نتایج", lines=8)
        img_btn = gr.Button("تحلیل عکس")
        img_btn.click(fn=analyze_image, inputs=img_input, outputs=img_output)

demo.launch()