import numpy as np
import pandas as pd
from PIL import Image
import streamlit as st
from streamlit_drawable_canvas import st_canvas
from streamlit_image_select import image_select
from streamlit_sortables import sort_items
from py4matching import template as m

def expand2square(imgpath, background_color=(0, 0, 0)):
    pil_img = Image.open(imgpath)
    width, height = pil_img.size
    if width == height:
        return pil_img
    elif width > height:
        result = Image.new(pil_img.mode, (width, width), background_color)
        result.paste(pil_img, (0, (width - height) // 2))
        return result.resize((700, 700))
    else:
        result = Image.new(pil_img.mode, (height, height), background_color)
        result.paste(pil_img, ((height - width) // 2, 0))
        return result.resize((700, 700))


@st.cache_data
def loading_data(files):
    imgs = []
    imgs_names = []
    imgs_dict = {}
    for file in files:
        image = expand2square(file)
        imgs.append(image)
        imgs_names.append(file.name)
        imgs_dict[file.name] = image
    return imgs, imgs_names, imgs_dict

if 'uploaded' not in st.session_state:
    st.session_state['uploaded'] = False

images = st.sidebar.file_uploader("Upload here the images (max 4 imgs for demo version):",
                                    type=["png", "jpg"], accept_multiple_files=True)


if len(images) > 0:
    st.session_state['uploaded'] = True
    imgs_path = []
    imgs = []
else:
    st.session_state['uploaded'] = False

if st.session_state['uploaded'] is True:

    # Loading uploaded images and cache the data
    imgs, imgs_path, imgs_dict = loading_data(images)

    # Specify canvas parameters in application
    drawing_mode = st.sidebar.selectbox(
        "Drawing tool:", ("point", "line", "rect", "circle", "transform")
    )

    stroke_width = st.sidebar.slider("Stroke width: ", 1, 25, 3)
    if drawing_mode == 'point':
        point_display_radius = st.sidebar.slider("Point display radius: ", 1, 25, 3)
    stroke_color = st.sidebar.color_picker("Stroke color hex: ")
    bg_color = st.sidebar.color_picker("Background color hex: ", "#eee")

    realtime_update = st.sidebar.checkbox("Update in realtime", False)

    master_index = image_select("Uploaded images", imgs, captions=imgs_path, return_value="index")

    # Create a canvas component
    canvas_result = st_canvas(
        fill_color="rgba(255, 165, 0, 0.3)",  # Fixed fill color with some opacity
        stroke_width=stroke_width,
        stroke_color=stroke_color,
        background_color=bg_color,
        background_image= imgs[master_index], #expand2square(bg_image) if bg_image else expand2square("./IMG_02099.jpg"),
        update_streamlit=realtime_update,
        height=700,
        width=700,
        drawing_mode=drawing_mode,
        point_display_radius=point_display_radius if drawing_mode == 'point' else 0,
        key="canvas",
    )

    test = st.sidebar.write("Select the processing order of slave images")

    with st.sidebar:
        imgs_path2 = imgs_path.copy()
        imgs_path2.pop(master_index)
        sorted_items = sort_items(imgs_path2, multi_containers=False, direction='vertical')

    TEMPLATE_SIZE = 64
    SEARCH_BUFFER = 32


    # if canvas_result.image_data is not None:
    #    st.image(canvas_result.image_data)
    if canvas_result.json_data is not None:
        df = pd.json_normalize(canvas_result.json_data["objects"])  # need to convert obj to str because PyArrow

        if 'type' in df.keys():
            #st.write(df.keys())
            df = df[['type', 'left', 'top']]
            df[sorted_items] = np.nan

            for index, row in df.iterrows():
                #print(row['type'], row['top'])
                i = int(row['left'])
                j = int(row['top'])

                template = np.array(imgs[master_index])[j:j+TEMPLATE_SIZE,
                                                        i:i+TEMPLATE_SIZE]
                # Loop on slave imgs
                st.image(template)
                for item in sorted_items:
                    st.write(item)

                    search_area = np.array(imgs_dict[item])[j-SEARCH_BUFFER:j+TEMPLATE_SIZE+SEARCH_BUFFER,
                                                            i-SEARCH_BUFFER:i+TEMPLATE_SIZE+SEARCH_BUFFER]
                    st.image(search_area)

                    px, py, max_val = m.template_match(template.astype('uint8'), search_area.astype('uint8'))
                    st.write(py)

            #Covert df to string for printing
            for col in df.select_dtypes(include=['object']).columns:
                df[col] = df[col].astype("str")
            st.dataframe(df)