import streamlit as st
from torchvision.transforms import functional as F
import gc
import numpy as np

from modules.streamlit_utils import *


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 uploaded_file is not None:
        get_score_threshold(is_mobile)
    
        if st.button("🚀 Launch Prediction"):
            st.session_state.image = 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()

    # Create placeholders for all sections
    prediction_result_placeholder = st.empty()
    additional_options_placeholder = st.empty()
    modeler_placeholder = st.empty()


    if 'prediction' in st.session_state and uploaded_file:
        if st.session_state.image != cropped_image:
            print('Image has changed')
            # Delete the prediction
            del st.session_state.prediction
            return

        with prediction_result_placeholder.container():
            if is_mobile:
                display_options(st.session_state.crop_image, st.session_state.score_threshold, is_mobile, int(5/6*screen_width))
            else:
                with st.expander("Show result of prediction"):
                    display_options(st.session_state.crop_image, st.session_state.score_threshold, is_mobile, int(5/6*screen_width))

        if not is_mobile:
            with additional_options_placeholder.container():
                modify_results()

        with modeler_placeholder.container():
            modeler_options(is_mobile)
            display_bpmn_modeler(is_mobile, screen_width)
    else:
        prediction_result_placeholder.empty()
        additional_options_placeholder.empty()
        modeler_placeholder.empty()
        # Create a lot of space for scrolling
        for _ in range(50):
            st.text("")

    gc.collect()

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