import streamlit as st
from transformers import pipeline
from rembg import remove
from PIL import Image
from io import BytesIO
import base64
from transformers import GPT2LMHeadModel, GPT2Tokenizer
import random

st.set_page_config(layout="wide", page_title="Image Classification App")
st.write("## Image Food Classification App")
st.sidebar.write("## Upload and download :gear:")

# Initialize image classification and recipe generation models
image_classifier = pipeline("image-classification", model="mjsp/sweet")
recipe_model = GPT2LMHeadModel.from_pretrained("gpt2")
recipe_tokenizer = GPT2Tokenizer.from_pretrained("gpt2")

MAX_FILE_SIZE = 5 * 1024 * 1024  # 5MB

def convert_image(img):
    buf = BytesIO()
    img.save(buf, format="PNG")
    byte_im = buf.getvalue()
    return byte_im

def fix_image(upload):
    image = Image.open(upload)
    col1.write("Original Image :camera:")
    col1.image(image)

    fixed = remove(image)
    col2.write("Fixed Image :wrench:")
    col2.image(fixed)
    st.sidebar.markdown("\n")
    st.sidebar.download_button("Download fixed image", convert_image(fixed), "fixed.png", "image/png")

def generate_recipe(title, max_length=200):
    matching_entries = [entry for entry in dataset if entry["title"] == title]

    if matching_entries:
        selected_entry = random.choice(matching_entries)
        title = selected_entry["title"]
        ingredients = selected_entry["ingredients"]
    else:
        title = "Default Recipe Title"
        ingredients = []

    input_text = f"Title: {title}\nIngredients: {', '.join(ingredients)}\n Instructions:"
    input_ids = recipe_tokenizer.encode(input_text, return_tensors="pt")

    output = recipe_model.generate(input_ids, max_length=max_length, num_return_sequences=1)

    generated_recipe = recipe_tokenizer.decode(output[0], skip_special_tokens=True)
    return generated_recipe

# Your dataset of titles, ingredients, and recipes
dataset = [
    {
        "title": "Gulab Jamun",
        "ingredients": ["milk powder", "ghee", "rose water", "saffron", "cardamom", "sugar syrup"],
        "recipe": "Instructions for making Gulab Jamun..."
    },
    {
        "title": "Jalebi",
        "ingredients": ["all-purpose flour", "yogurt", "sugar", "water", "saffron strands", "cardamom powder", "ghee or oil for frying"],
        "recipe": "Instructions for making Jalebi..."
    },
    {
        "title": "Rasgulla",
        "ingredients": ["milk", "sugar", "lemon juice", "rose water"],
        "recipe": "Instructions for making Rasgulla..."
    }
]

col1, col2 = st.columns(2)
my_upload = st.sidebar.file_uploader("Upload an image", type=["png", "jpg", "jpeg"])

st.write()
st.markdown("""---""")
st.write()

# global label
# label = ""

if my_upload is not None:
    st.image(my_upload, caption="Uploaded Image", use_column_width=True)

    if st.sidebar.button("Classify"):
        st.sidebar.text("Classifying...")

        image = Image.open(my_upload)

        try:
            classification_result = image_classifier(image)
            top_prediction = classification_result[0]
            label = top_prediction['label']
            score = top_prediction['score']
            st.sidebar.text("Top Prediction:")
            st.sidebar.text(f"Label: {label}, Score: {score:.3f}")
        except Exception as e:
            st.error(f"Error during classification: {str(e)}")

    if my_upload.size > MAX_FILE_SIZE:
        st.error("The uploaded file is too large. Please upload an image smaller than 5MB.")
    else:
        fix_image(upload=my_upload)
else:
    fix_image("jalebi.jpg")



# Recipe generation based on selected item
st.write("## Recipe Generation")

selected_item = st.selectbox("Select a food item", [entry["title"] for entry in dataset])
if st.button("Generate Recipe"):
     generated_recipe = generate_recipe(selected_item, max_length=200)
     st.write(f"Recipe for {selected_item}:\n{generated_recipe}")

# Add some descriptions and instructions
st.sidebar.markdown("### Instructions")
st.sidebar.markdown("1. Upload an image.")
st.sidebar.markdown("2. Click the 'Classify' button to get the classification results.")
st.sidebar.markdown("3. Select a food item to generate a recipe.")
st.sidebar.markdown("4. Click the 'Generate Recipe' button to get the recipe.")

# Display a placeholder for the main content
st.write("Please upload an image and use the sidebar to classify it and generate a recipe.")