import streamlit as st
from streamlit_image_select import image_select
from PIL import Image
import numpy as np

# UI configurations
st.set_page_config(page_title="AIFR - Demo",
                   page_icon=":bridge_at_night:",
                   layout="wide")
st.markdown("# :rainbow[AIFR - Demo]")

# 3 columns
col1, col2, col3, col4 = st.columns(4)

with col1:
    st.header("User Image")

    user_image_holder = st.empty()
    # upload file
    user_image = st.file_uploader("Upload User Image")
    if user_image is not None:
        img = None
        user_image_holder.image(user_image, use_column_width=True)
    
    st.write("Examples")
    img1 = image_select(
        label="Select a cat",
        images=[
            "https://bagongkia.github.io/react-image-picker/0759b6e526e3c6d72569894e58329d89.jpg",
            "https://bagongkia.github.io/react-image-picker/0759b6e526e3c6d72569894e58329d89.jpg"
        ],
        captions=["A cat", "Another cat"],
    )
    if img1 and user_image is None:
        user_image = img1
        user_image_holder.image(user_image, use_column_width=True)

with col2:
    st.header("Clothes Image")

    clothes_image_holder = st.empty()
    # upload file
    clothes_image = st.file_uploader("Upload Clothes Image")
    if clothes_image is not None:
        clothes_image_holder.image(clothes_image, use_column_width=True)

    st.write("Examples")
    img2 = image_select(
        label="Select a dress",
        images=[
            "https://bagongkia.github.io/react-image-picker/0759b6e526e3c6d72569894e58329d89.jpg",
            "https://bagongkia.github.io/react-image-picker/0759b6e526e3c6d72569894e58329d89.jpg"
        ],
        captions=["A dress", "Another dress"],
    )

    if img2 and clothes_image is None:
        clothes_image = img2
        clothes_image_holder.image(clothes_image, use_column_width=True)


with col3:
    st.header("Masked Image output")

    mask_image = None
    if mask_image is not None:
        st.image(mask_image, use_column_width=True)

with col4:
    st.header("Output")

    result_image = None
    if result_image is not None:
        st.image(result_image, use_column_width=True)