import streamlit as st
from torchvision.transforms import functional as F
import gc
import numpy as np
from modules.htlm_webpage import display_bpmn_xml
from streamlit_cropper import st_cropper
from streamlit_image_select import image_select
from streamlit_js_eval import streamlit_js_eval
from streamlit_drawable_canvas import st_canvas
from modules.streamlit_utils import *
from glob import glob

from streamlit_image_annotation import detection
from modules.toXML import create_XML
from modules.eval import develop_prediction, generate_data
from modules.utils import class_dict, object_dict

def configure_page():
    st.set_page_config(layout="wide")
    screen_width = streamlit_js_eval(js_expressions='screen.width', want_output=True, key='SCR')
    is_mobile = screen_width is not None and screen_width < 800
    return is_mobile, screen_width

def display_banner(is_mobile):
    if is_mobile:
        st.image("./images/banner_mobile.png", use_column_width=True)
    else:
        st.image("./images/banner_desktop.png", use_column_width=True)

def display_title(is_mobile):
    title = "Welcome on the BPMN AI model recognition app"
    if is_mobile:
        title = "Welcome on the mobile version of BPMN AI model recognition app"
    st.title(title)

def display_sidebar():
    sidebar()

def initialize_session_state():
    if 'pool_bboxes' not in st.session_state:
        st.session_state.pool_bboxes = []
    if 'model_object' not in st.session_state or 'model_arrow' not in st.session_state:
        clear_memory()
        load_models()

def load_example_image():
    with st.expander("Use example images"):
        img_selected = image_select(
            "If you have no image and just want to test the demo, click on one of these images", 
            ["./images/none.jpg", "./images/example1.jpg", "./images/example2.jpg", "./images/example3.jpg", "./images/example4.jpg"],
            captions=["None", "Example 1", "Example 2", "Example 3", "Example 4"], 
            index=0, 
            use_container_width=False, 
            return_value="original"
        )
        return img_selected

def load_user_image(img_selected, is_mobile):
    if img_selected == './images/none.jpg':
        img_selected = None

    if img_selected is not None:
        uploaded_file = img_selected
    else:
        if is_mobile:
            uploaded_file = st.file_uploader("Choose an image from my computer...", type=["jpg", "jpeg", "png"], accept_multiple_files=False)
        else:
            col1, col2 = st.columns(2)
            with col1:
                uploaded_file = st.file_uploader("Choose an image from my computer...", type=["jpg", "jpeg", "png"])

    return uploaded_file

def display_image(uploaded_file, screen_width, is_mobile):
    
    with st.spinner('Waiting for image display...'):
        original_image = get_image(uploaded_file)
        resized_image = original_image.resize((screen_width // 2, int(original_image.height * (screen_width // 2) / original_image.width)))

        if not is_mobile:
            cropped_image = crop_image(resized_image, original_image)
        else:
            st.image(resized_image, caption="Image", use_column_width=False, width=int(4/5 * screen_width))
            cropped_image = original_image

    return cropped_image

def crop_image(resized_image, original_image):
    marge = 10
    cropped_box = st_cropper(
        resized_image,
        realtime_update=True,
        box_color='#0000FF',
        return_type='box',
        should_resize_image=False,
        default_coords=(marge, resized_image.width - marge, marge, resized_image.height - marge)
    )
    scale_x = original_image.width / resized_image.width
    scale_y = original_image.height / resized_image.height
    x0, y0, x1, y1 = int(cropped_box['left'] * scale_x), int(cropped_box['top'] * scale_y), int((cropped_box['left'] + cropped_box['width']) * scale_x), int((cropped_box['top'] + cropped_box['height']) * scale_y)
    cropped_image = original_image.crop((x0, y0, x1, y1))
    return cropped_image

def get_score_threshold(is_mobile):
    col1, col2 = st.columns(2)
    with col1:
        st.session_state.score_threshold = st.slider("Set score threshold for prediction", min_value=0.0, max_value=1.0, value=0.5 if not is_mobile else 0.6, step=0.05) 

def launch_prediction(cropped_image, score_threshold, is_mobile, screen_width):
    st.session_state.crop_image = cropped_image
    with st.spinner('Processing...'):
        perform_inference(
            st.session_state.model_object, st.session_state.model_arrow, st.session_state.crop_image,
            score_threshold, is_mobile, screen_width, iou_threshold=0.3, distance_treshold=30, percentage_text_dist_thresh=0.5
        )
        st.balloons()    
    
def mix_new_pred(objects_pred, arrow_pred):
    # Initialize the list of lists for keypoints
    object_keypoints = []

    # Number of boxes
    num_boxes = len(objects_pred['boxes'])

    # Iterate over the number of boxes
    for _ in range(num_boxes):
        # Each box has 2 keypoints, both initialized to [0, 0, 0]
        keypoints = [[0, 0, 0], [0, 0, 0]]
        object_keypoints.append(keypoints)

    #concatenate the two predictions
    boxes = np.concatenate((objects_pred['boxes'], arrow_pred['boxes']))
    labels = np.concatenate((objects_pred['labels'], arrow_pred['labels']))

    return boxes, labels, keypoints


def modify_results(percentage_text_dist_thresh=0.5):
    with st.expander("Method and Style modification (beta version)"):
        label_list = list(object_dict.values())
        bboxes = [[int(coord) for coord in box] for box in st.session_state.prediction['boxes']]
        for i in range(len(bboxes)):
            bboxes[i][2] = bboxes[i][2] - bboxes[i][0]
            bboxes[i][3] = bboxes[i][3] - bboxes[i][1]
        labels = [int(label) for label in st.session_state.prediction['labels']]


        # Filter boxes and labels where label is less than 12
        ignore_labels = [6, 7]
        object_bboxes = []
        object_labels = []
        arrow_bboxes = []
        arrow_labels = []
        for i in range(len(bboxes)):
            if labels[i] <= 12:
                object_bboxes.append(bboxes[i])
                object_labels.append(labels[i])
            else:
                arrow_bboxes.append(bboxes[i])
                arrow_labels.append(labels[i])

        print('Object bboxes:', object_bboxes)
        print('Object labels:', object_labels)
        print('Arrow bboxes:', arrow_bboxes)
        print('Arrow labels:', arrow_labels)

        original_obj_len = len(object_bboxes)


        uploaded_image = prepare_image(st.session_state.crop_image, new_size=(1333, 1333), pad=False)
        scale = 2000 / uploaded_image.size[0]
        new_labels = detection(
            image=uploaded_image, bboxes=object_bboxes, labels=object_labels, 
            label_list=label_list, line_width=3, width=2000, use_space=False
        )

        if new_labels is not None:
            new_lab = np.array([label['label_id'] for label in new_labels])  
            # Convert back to original format
            bboxes = np.array([label['bbox'] for label in new_labels])
            for i in range(len(bboxes)):
                bboxes[i][2] = bboxes[i][2] + bboxes[i][0]
                bboxes[i][3] = bboxes[i][3] + bboxes[i][1]
            for i in range(len(arrow_bboxes)):
                arrow_bboxes[i][2] = arrow_bboxes[i][2] + arrow_bboxes[i][0]
                arrow_bboxes[i][3] = arrow_bboxes[i][3] + arrow_bboxes[i][1]


            new_bbox = np.concatenate((bboxes, arrow_bboxes))
            new_lab = np.concatenate((new_lab, arrow_labels))

            print('New labels:', new_lab)

            scores = st.session_state.prediction['scores']
            keypoints = st.session_state.prediction['keypoints']

            #delete element in keypoints to make it match the new number of boxes
            len_keypoints = len(keypoints)
            keypoints = keypoints.tolist()
            scores = scores.tolist()


            diff = original_obj_len-len(bboxes)
            if diff > 0:
                for i in range(diff):                  
                            keypoints.pop(0)
                            scores.pop(0)
            elif diff < 0:
                for i in range(-diff):
                    keypoints.insert(0, [[0, 0, 0], [0, 0, 0]])
                    scores.insert(0, 0.0)

            print('lenghts: ',len(bboxes), len(new_lab), len(scores), len(keypoints))
            keypoints = np.array(keypoints)
            scores = np.array(scores)

            #print('Old prediction:', st.session_state.prediction['keypoints'])
            boxes, labels, scores, keypoints, flow_links, best_points, pool_dict = develop_prediction(new_bbox, new_lab, scores, keypoints, class_dict, correction=False)

            st.session_state.prediction = generate_data(st.session_state.prediction['image'], boxes, labels, scores, keypoints, flow_links, best_points, pool_dict, class_dict)
            st.session_state.text_mapping = mapping_text(st.session_state.prediction, st.session_state.text_pred, print_sentences=False, percentage_thresh=percentage_text_dist_thresh)

            #print('New prediction:', st.session_state.prediction['keypoints'])
            st.rerun()


def display_bpmn_modeler(is_mobile, screen_width):
    with st.spinner('Waiting for BPMN modeler...'):
        st.session_state.bpmn_xml = create_XML(
            st.session_state.prediction.copy(), st.session_state.text_mapping, 
            st.session_state.size_scale, st.session_state.scale
        )
        display_bpmn_xml(st.session_state.bpmn_xml, is_mobile=is_mobile, screen_width=int(4/5 * screen_width))

def modeler_options(is_mobile):
    if not is_mobile:
        with st.expander("Options for BPMN modeler"):
            col1, col2 = st.columns(2)
            with col1:
                st.session_state.scale = st.slider("Set distance scale for XML file", min_value=0.1, max_value=2.0, value=1.0, step=0.1) 
                st.session_state.size_scale = st.slider("Set size object scale for XML file", min_value=0.5, max_value=2.0, value=1.0, step=0.1) 
    else:
        st.session_state.scale = 1.0
        st.session_state.size_scale = 1.0

def main():
    is_mobile, screen_width = configure_page()
    display_banner(is_mobile)
    display_title(is_mobile)
    display_sidebar()
    initialize_session_state()

    cropped_image = None

    img_selected = load_example_image()
    uploaded_file = load_user_image(img_selected, is_mobile)
    if uploaded_file is not None:
        cropped_image = display_image(uploaded_file, screen_width, is_mobile)

    if cropped_image is not None:
        get_score_threshold(is_mobile)
        if st.button("🚀 Launch Prediction"):
            launch_prediction(cropped_image, st.session_state.score_threshold, is_mobile, screen_width)
            st.session_state.original_prediction = st.session_state.prediction.copy()
            st.rerun()
    
    if 'prediction' in st.session_state and uploaded_file:
        #if st.button("🔄 Refresh image"):
            #st.rerun()

        with st.expander("Show result of prediction"):
            with st.spinner('Waiting for result display...'):
                display_options(st.session_state.crop_image, st.session_state.score_threshold, is_mobile, int(5/6 * screen_width))

        if not is_mobile:
            modify_results()

        modeler_options(is_mobile)
                
        
        display_bpmn_modeler(is_mobile, screen_width)

    gc.collect()

if __name__ == "__main__":
    print('Starting the app...')
    main()