import streamlit as st
import tensorflow as tf
from tensorflow.keras.preprocessing import image
import numpy as np
import os

# Load the model
MODEL_PATH = "./final_model.keras"
model = tf.keras.models.load_model(MODEL_PATH)

# Define the class names directly for the Food-101 dataset
CLASS_NAMES = [
    'apple_pie', 'baby_back_ribs', 'baklava', 'beef_carpaccio', 'beef_tartare', 'beet_salad', 'beignets',
    'bibimbap', 'bread_pudding', 'breakfast_burrito', 'bruschetta', 'caesar_salad', 'cannoli', 'caprese_salad',
    'carrot_cake', 'ceviche', 'cheese_plate', 'cheesecake', 'chicken_curry', 'chicken_quesadilla',
    'chicken_wings', 'chocolate_cake', 'chocolate_mousse', 'churros', 'clam_chowder', 'club_sandwich',
    'crab_cakes', 'creme_brulee', 'croque_madame', 'cup_cakes', 'deviled_eggs', 'donuts', 'dumplings',
    'edamame', 'eggs_benedict', 'escargots', 'falafel', 'filet_mignon', 'fish_and_chips', 'foie_gras',
    'french_fries', 'french_onion_soup', 'french_toast', 'fried_calamari', 'fried_rice', 'frozen_yogurt',
    'garlic_bread', 'gnocchi', 'greek_salad', 'grilled_cheese_sandwich', 'grilled_salmon', 'guacamole',
    'gyoza', 'hamburger', 'hot_and_sour_soup', 'hot_dog', 'huevos_rancheros', 'hummus', 'ice_cream',
    'lasagna', 'lobster_bisque', 'lobster_roll_sandwich', 'macaroni_and_cheese', 'macarons', 'miso_soup',
    'mussels', 'nachos', 'omelette', 'onion_rings', 'oysters', 'pad_thai', 'paella', 'pancakes', 'panna_cotta',
    'peking_duck', 'pho', 'pizza', 'pork_chop', 'poutine', 'prime_rib', 'pulled_pork_sandwich', 'ramen',
    'ravioli', 'red_velvet_cake', 'risotto', 'samosa', 'sashimi', 'scallops', 'seaweed_salad', 'shrimp_and_grits',
    'spaghetti_bolognese', 'spaghetti_carbonara', 'spring_rolls', 'steak', 'strawberry_shortcake', 'sushi',
    'tacos', 'takoyaki', 'tiramisu', 'tuna_tartare', 'waffles'
]

# Define the function to predict the image
def predict_image(img_path):
    # Load and preprocess the image
    img = image.load_img(img_path, target_size=(224, 224))  # Resize to match model's expected input size
    img_array = image.img_to_array(img)  # Convert image to array
    img_array = np.expand_dims(img_array, axis=0)  # Add batch dimension
    img_array = img_array / 255.0  # Normalize the image as done during training

    # Make prediction
    predictions = model.predict(img_array)

    # Get prediction probabilities
    prediction_probs = predictions[0]  # Prediction probabilities
    predicted_class_index = np.argmax(prediction_probs)
    predicted_class = CLASS_NAMES[predicted_class_index]  # Fetch the class name

    return predicted_class


# Streamlit UI components
st.set_page_config(page_title="Food-101 Classification", page_icon="🍕", layout="wide")
st.markdown("""
    <style>
    .stApp {
        background-color: #f03c45;
    }
    .stButton>button {
        background-color: #4CAF50;
        color: white;
        padding: 10px 20px;
        border-radius: 10px;
        font-size: 16px;
    }
    .stFileUploader>label {
        color: #4CAF50;
    }
    .stImage>img {
        border-radius: 10px;
        box-shadow: 0px 4px 12px rgba(0, 0, 0, 0.1);
    }
    </style>
""", unsafe_allow_html=True)

# Title
st.title("🍕 Food-101 Classification Model")
st.subheader("Upload an image of food to predict its class. 🍽️")

# Upload image
uploaded_file = st.file_uploader("Choose an image...", type=["jpg", "jpeg", "png"])

if uploaded_file is not None:
    # Display the uploaded image with a transition effect
    st.image(uploaded_file, caption="Uploaded Image", use_column_width=True, width=500)

    # Save the image temporarily
    img_path = "uploaded_image.jpg"
    with open(img_path, "wb") as f:
        f.write(uploaded_file.getbuffer())

    # Show the processing animation
    with st.spinner("Classifying the image... Please wait. 🕒"):
        # Make prediction
        predicted_class = predict_image(img_path)

    # Display the predicted class with a colorful box
    st.markdown(f"""
        <div style="padding: 10px; border-radius: 10px; background-color: #ebe7d3; color: #0f4b5f; text-align: center; font-size: 20px;">
            <strong>Predicted Class: {predicted_class.capitalize()} 🥳</strong>
        </div>
    """, unsafe_allow_html=True)

    # Show an icon with a success message
    # st.balloons()