import streamlit as st
from PIL import Image
import cv2
from texture_transfer import create_image_tile, create_layover, create_image_cutout, paste_image
from create_print_layover import create_image_with_feather_tile, control_texture

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")
    choice = st.selectbox(
        'what are you looking for?',
        ("Direct Texture", "Feathered Texture", "Direction Print"),
        index=None,
        placeholder="please choose an option"
    )
    overlap_width = st.slider("Overlap Width", min_value=0, max_value=100, step=1, value=50)
    what_overlap = st.selectbox(
        'how you want to overlap?',
        ('vertical', 'horizontal'),
        index=None,
        placeholder="please choose an option"
    )
    opacity = st.slider("Control opacity", min_value=0.0, max_value=1.0, value=1.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_n = Image.open(generated_image)
            product_image_dim = product_image_n.size
            x_dim = product_image_dim[0]
            y_dim = product_image_dim[1]
            if choice == "Direct Texture":
                create_image_tile(product_image, x_dim, y_dim)
            elif choice == "Feathered Texture":
                create_image_with_feather_tile(product_image, x_dim, y_dim, overlap_width)
            elif choice == "Direction Print":
                control_texture(product_image, what_overlap, overlap_width, x_dim, y_dim)
                color_texture_image = cv2.imread('tiled_image_2.png')
                resized_image = cv2.resize(color_texture_image, (x_dim, y_dim))
                cv2.imwrite('tiled_image.png', resized_image)
            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("tiled_image.png", caption="tiled_image.png")
            # st.image("lay_over_image.png", caption="lay_over_image.png")
            # st.image("cut_out_image.png", caption="cut_out_image.png")
            st.image("result.png", caption="Texture Corrected Image", use_column_width=True)