import streamlit as st
from PIL import Image
from texture_transfer import create_image_tile, create_layover, create_image_cutout, paste_image

st.title("Texture Correction App")

if 'clicked' not in st.session_state:
    st.session_state.clicked = False


def click_button():
    st.session_state.clicked = True


with st.sidebar:
    st.text("Play around this values")
    opacity = st.slider("Control opacity", min_value=0.0, max_value=100.0, value=100.0)

product, generated, texture_correct = st.columns(3)
with product:
    st.header("Input product texture patch")
    product_image = st.file_uploader("Patch of the Garment(Cut out an area from the garment)",
                                     type=["png", "jpg", "jpeg"])
    if product_image:
        st.image(product_image)

with generated:
    st.header("Input generated product Image & Mask")
    generated_image = st.file_uploader("Generated Image", type=["png", "jpg", "jpeg"])
    gen_image_mask = st.file_uploader("Generated Mask", type=["png", "jpg", "jpeg"])
    if generated_image:
        st.image(generated_image)

with texture_correct:
    st.button("Texture Correct", on_click=click_button)
    if st.session_state.clicked:
        with st.spinner("Texture correcting🫡..."):
            product_image = Image.open(generated_image)
            product_image_dim = product_image.size
            x_dim = product_image_dim[0]
            y_dim = product_image_dim[0]
            create_image_tile(product_image, x_dim, y_dim)
            overlay = create_layover(generated_image, 'tiled_image.png', opacity)
            create_image_cutout('lay_over_image.png', gen_image_mask)
            paste_image(generated_image, 'cut_out_image.png', gen_image_mask)
            st.image("result.png", caption="Texture Corrected Image", use_column_width=True)