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 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()