import streamlit as st
import requests
import tempfile
import numpy as np

# Set page config
st.set_page_config(page_title='Image Generation', page_icon=':pizza:', layout='wide')

st.title('Pizza Image Generation')

# Initialize session state for storing user inputs
if 'dynamic_prompt' not in st.session_state:
    st.session_state['dynamic_prompt'] = ""
if 'image_data' not in st.session_state:
    st.session_state['image_data'] = None

# Define pizza options
pizza_names = ["Crispy Chicken", "Chicken Pesto", "Deluxe Pepperoni", "Truffle Mushroom",
               "Ultimate Cheese", "Spicy Veggie Ranch", "Classic Chicken Ranch",
               "BBQ Chicken Ranch", "Spicy Chicken Ranch", "Very Veggie",
               "Super Supreme", "Classic Pepperoni", "Margherita", "Cheeky Chicken",
               "Chicken Super Supreme", "Chicken BBQ Supreme", "Chicken Fajita",
               "Chicken Shawerma", "Other"]

pizza_sizes = ["Small", "Medium", "Large", "Extra Large"]

toppings_options = ["Arugula", "Bacon", "Basil", "Broccoli", "Cheese", "Chicken", "Corn", 
                    "Ham", "Mushroom", "Olives", "Onion", "Pepperoni", "Peppers", 
                    "Pineapple", "Tomatoes"]

bg_color_options = ['White', 'Black', 'Light Blue', 'Light Green', 'Light Yellow', 'Light Pink']

orientation_options = ["Top View", "Side View", "Lateral View"]

# UI Components for Pizza Preferences
with st.form(key='prompt_form'):
    # Pizza selection
    selected_pizza = st.selectbox("Choose a Pizza Type", pizza_names)
    selected_size = st.selectbox("Choose Pizza Size", pizza_sizes)
    
    # Toppings selection
    selected_toppings = st.multiselect("Choose Toppings", toppings_options, default=["Mushroom", "Tomatoes", "Onion"])

    # Additional customization
    additional_description = st.text_input("Enter additional descriptions (optional)")
    background_color = st.selectbox("Choose Background Color", bg_color_options)
    orientation = st.multiselect("Select Orientation", orientation_options, default=["Top View"])

    generate_prompt_button = st.form_submit_button(label='Generate Prompt')

# Processing user input to generate dynamic prompt
if generate_prompt_button:
    st.session_state['dynamic_prompt'] = f"A {selected_size} {selected_pizza} Pizza with {', '.join(selected_toppings)}. Additional description: {additional_description}. Background color: {background_color}. Orientation: {', '.join(orientation)}."
    st.write("Prompt:", st.session_state['dynamic_prompt'])

# Rest of your code for image generation and display...
with st.form(key='image_form'):
    generate_image_button = st.form_submit_button(label='Generate Image')

if generate_image_button and st.session_state['dynamic_prompt']:
    # Generate the image using your custom model
    # Assuming 'my_custom_model_generate_image' is a function that returns the image data
    st.session_state['image_data'] = (st.session_state['dynamic_prompt'])

    # Display the generated image if available
    if st.session_state['image_data'] is not None:
        st.image(st.session_state['image_data'], caption="Generated Image", use_column_width=True)