import streamlit as st
from transformers import pipeline
from PIL import Image
from io import BytesIO
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)
    st.image(image, caption="Original Image", use_column_width=True)

    # You'll need to add the 'rembg' functionality or replace it with your own image processing logic
    # fixed = your_image_processing_function(image)
    # st.image(fixed, caption="Fixed Image", use_column_width=True)
    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):
    # Replace this with your actual dataset
    dataset = {
        "Gulab Jamun": {
            "ingredients": ["milk powder", "ghee", "rose water", "saffron", "cardamom", "sugar syrup"],
            "recipe": "Instructions for making Gulab Jamun...",
        },
        "Jalebi": {
            "ingredients": ["all-purpose flour", "yogurt", "sugar", "water", "saffron strands", "cardamom powder", "ghee or oil for frying"],
            "recipe": "Instructions for making Jalebi...",
        },
        "Rasgulla": {
            "ingredients": ["milk", "sugar", "lemon juice", "rose water"],
            "recipe": "Instructions for making Rasgulla...",
        }
    }

    if title in dataset:
        selected_entry = dataset[title]
        title = 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

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

if my_upload:
    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(my_upload)

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

selected_item = st.selectbox("Select a food item", ["Gulab Jamun", "Jalebi", "Rasgulla"])
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.")