import streamlit as st
import tensorflow as tf
import numpy as np
from PIL import Image
import pandas as pd
import matplotlib.pyplot as plt

# Load the saved model 
model_path = "fahrnphi_exam_project.keras"
model = tf.keras.models.load_model(model_path)

# Define the core prediction function
def predict_ingredient(image):
    # Preprocess image
    image = image.resize((150, 150))  # Resize the image to 150x150
    image = image.convert('RGB')  # Ensure image has 3 channels
    image = np.array(image)
    image = np.expand_dims(image, axis=0)  # Add batch dimension
    
    # Predict
    prediction = model.predict(image)
    
    # Apply softmax to get probabilities for each class
    probabilities = tf.nn.softmax(prediction, axis=1)
    
    # Map probabilities to ingredient classes
    class_names = ['Peperoni', 'Carrot', 'Garlic', 'Ginger', 'Jalapeno', 'Onion', 'Potato', 'Sweetpotato', 'Tomato']
    probabilities_dict = {ingredient_class: round(float(probability), 2) for ingredient_class, probability in zip(class_names, probabilities.numpy()[0])}
    
    return probabilities_dict

# Streamlit interface
st.title("Ingredient Classifier")
st.write("A simple MLP classification model for image classification using a pretrained model.")

# Initialize session state for storing ingredients
if 'ingredients' not in st.session_state:
    st.session_state['ingredients'] = []

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

if uploaded_image is not None:
    image = Image.open(uploaded_image)
    st.image(image, caption='Uploaded Image.', use_column_width=True)
    st.write("")
    st.write("Classifying...")
    
    predictions = predict_ingredient(image)

    # Display predictions as a DataFrame
    st.write("### Prediction Probabilities")
    df = pd.DataFrame(predictions.items(), columns=["Ingredient", "Probability"])
    st.dataframe(df)
    
    # Display predictions as a pie chart
    st.write("### Prediction Chart")
    fig, ax = plt.subplots()
    ax.pie(df["Probability"], labels=df["Ingredient"], autopct='%1.1f%%', colors=plt.cm.Paired.colors)
    ax.set_title('Prediction Probabilities')
    st.pyplot(fig)
    
    # Automatically select the best guess (highest probability)
    best_guess = df.loc[df['Probability'].idxmax()]["Ingredient"]
    
    if st.button("Add Ingredient"):
        st.session_state.ingredients.append(best_guess)
        st.write(f"Added {best_guess} to ingredients list")

# Display the ingredients added so far
st.write("### Selected Ingredients")
st.write(st.session_state.ingredients)

# Finish button to finalize ingredient selection and find recipes
if st.button("Finish"):
    st.write("Finding recipes...")
    # Placeholder: Replace with actual recipe finding logic
    def find_recipes(ingredients):
        # This is a mock function, replace with actual recipe finding logic
        sample_recipes = [
            {"name": "Vegetable Stir Fry", "ingredients": ["Peperoni", "Carrot", "Onion"], "instructions": "Stir fry the vegetables in a hot pan with some oil."},
            {"name": "Tomato Garlic Pasta", "ingredients": ["Tomato", "Garlic"], "instructions": "Cook pasta and mix with sautéed tomato and garlic."},
            {"name": "Ginger Potato Soup", "ingredients": ["Ginger", "Potato"], "instructions": "Boil potatoes and ginger, then blend into a soup."},
            {"name": "Jalapeno Onion Salad", "ingredients": ["Jalapeno", "Onion"], "instructions": "Mix chopped jalapeno and onion with some lime juice."},
            {"name": "Sweetpotato Carrot Soup", "ingredients": ["Sweetpotato", "Carrot"], "instructions": "Boil sweetpotato and carrot, then blend into a soup."},
            {"name": "Garlic Mashed Potatoes", "ingredients": ["Garlic", "Potato"], "instructions": "Boil potatoes, mash them with roasted garlic and butter."},
            {"name": "Ginger Carrot Salad", "ingredients": ["Ginger", "Carrot"], "instructions": "Grate carrots and mix with finely chopped ginger and a vinaigrette."},
            {"name": "Pepperoni Pizza", "ingredients": ["Peperoni", "Tomato", "Onion"], "instructions": "Top pizza dough with tomato sauce, peperoni, and onion slices. Bake until crispy."},
            {"name": "Onion Soup", "ingredients": ["Onion"], "instructions": "Sauté onions until caramelized, then add broth and simmer."},
            {"name": "Tomato Salad", "ingredients": ["Tomato", "Onion"], "instructions": "Chop tomatoes and onions, mix with olive oil and vinegar."},
            {"name": "Carrot Ginger Soup", "ingredients": ["Carrot", "Ginger"], "instructions": "Boil carrots and ginger, blend into a creamy soup."},
            {"name": "Potato Jalapeno Gratin", "ingredients": ["Potato", "Jalapeno"], "instructions": "Layer sliced potatoes and jalapenos, bake with cream and cheese."},
            {"name": "Garlic Ginger Stir Fry", "ingredients": ["Garlic", "Ginger"], "instructions": "Stir fry garlic and ginger with your choice of vegetables."},
            {"name": "Roasted Peperoni", "ingredients": ["Peperoni"], "instructions": "Roast whole peperoni in the oven until charred."},
            {"name": "Sweetpotato Fries", "ingredients": ["Sweetpotato"], "instructions": "Cut sweetpotatoes into fries, season, and bake until crispy."},
            {"name": "Garlic Ginger Chicken", "ingredients": ["Garlic", "Ginger"], "instructions": "Marinate chicken with garlic and ginger, then bake or grill."},
            {"name": "Onion Rings", "ingredients": ["Onion"], "instructions": "Dip onion slices in batter and deep fry until golden."},
            {"name": "Tomato Basil Bruschetta", "ingredients": ["Tomato"], "instructions": "Top toasted bread with diced tomatoes, basil, and olive oil."},
            {"name": "Jalapeno Poppers", "ingredients": ["Jalapeno"], "instructions": "Stuff jalapenos with cheese, bread them, and bake or fry."},
            {"name": "Carrot Sweetpotato Mash", "ingredients": ["Carrot", "Sweetpotato"], "instructions": "Boil carrots and sweetpotatoes, then mash with butter and seasoning."}
]
        matching_recipes = [recipe for recipe in sample_recipes if all(item in recipe["ingredients"] for item in ingredients)]
        return matching_recipes
    
    matching_recipes = find_recipes(st.session_state.ingredients)
    
    if matching_recipes:
        st.write("### Matching Recipes")
        for recipe in matching_recipes:
            st.write(f"**{recipe['name']}**")
            st.write(", ".join(recipe["ingredients"]))
            st.write(f"Instructions: {recipe['instructions']}")

    else:
        st.write("No matching recipes found.")

# Reset button to start over
if st.button("Reset"):
    st.session_state.ingredients = []
    st.write("Ingredients list has been reset.")