import streamlit as st
from PIL import Image

# Set the page config to make the app use the full width of the browser
st.set_page_config(layout="wide")

# Initialize a session state to keep track of images
if 'images' not in st.session_state:
    st.session_state.images = []

# Function to reset images
def reset_images():
    # Clear the images list
    st.session_state.images.clear()
    # Reset the file uploader widget
    st.session_state.file_uploader = None

# Main app function
def main():
    st.title("Sock Image Uploader")

    # Image uploader
    uploaded_file = st.file_uploader("Take a picture of your sock and upload it here", type=["jpg", "png"], key="file_uploader")

    # If a file is uploaded, process and display it immediately
    if uploaded_file is not None:
        # Convert the file to an image
        image = Image.open(uploaded_file)
        st.session_state.images.append(image)

    # Display images in a grid
    cols_per_row = 4
    for i in range(0, len(st.session_state.images), cols_per_row):
        cols = st.columns(cols_per_row)
        for j in range(cols_per_row):
            if i + j < len(st.session_state.images):
                with cols[j]:
                    st.image(st.session_state.images[i + j], caption=f"Image {i + j + 1}", use_column_width=True)

    # Reset button
    if st.button('Reset and Start Over'):
        reset_images()

if __name__ == "__main__":
    main()