import streamlit as st
from streamlit_image_select import image_select
from utils_streamlit import get_result
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)
    
    img1 = image_select(
        label="Examples",
        images=[
            "./user_example/ex1.jpg",
            "./user_example/ex2.jpg",
        ],
        use_container_width=False,
    )
    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)

    img2 = image_select(
        label="Example",
        images=[
            "./garment_example/ex1.jpg",
            "./garment_example/ex2.jpg",
            "./garment_example/ex3.jpg",
            "./garment_example/ex4.jpg",
        ],
        use_container_width=False,
    )

    if img2 and clothes_image is None:
        clothes_image = img2
        clothes_image_holder.image(clothes_image, use_column_width=True)
body_part = st.selectbox(
    "Choose your body part",
    ("dresses", "upper_body", "lower_body"))
submitted = st.button("Get result", use_container_width=True, type="primary")
output_image = mask_image = None

if submitted:
    user_image = Image.open(user_image)
    clothes_image = Image.open(clothes_image)
    output_image, mask_image = get_result(user_image, clothes_image, body_part=body_part)

with col3:
    st.header("Masked Image output")
    if submitted:
        if mask_image is not None:
            st.image(mask_image, use_column_width=True)

with col4:
    st.header("Output")
    if submitted:
        if output_image is not None:
            st.image(output_image, use_column_width=True)