import streamlit as st

from beluga import load_model, process_emotions, generate_prompt
from emodeepface import check_image_rotation, process_photo

# add streamlit cache to prevent multiple reloads of beluga model
@st.cache_resource
def load_cached_model():
    return load_model()

# If the model and tokenizer aren't already in session state, load them
if 'model' not in st.session_state:
    loading_message = st.empty()
    loading_message.text("Loading model... Please wait.")
    
    # begin loading beluga model and tokenizer
    st.session_state.model, st.session_state.tokenizer = load_cached_model()

    # clear loading message
    loading_message.empty()

# title webpage
st.title("Affective Journaling Assistant")

# provide user instructions
st.write("""
Welcome to the Affective Journaling Assistant!

For a tailored journaling experience, we analyze your facial expressions to gauge your emotions.
To proceed:

1. Ensure the image is well-lit and of high quality.
2. Your face should be fully visible without obstructions (e.g., no sunglasses or hats).
3. By uploading, you acknowledge and consent to our data processing.

Let's get started!
""")

# request user image upload
file_name = st.file_uploader("Please upload your photo.")

# once an image has been uploaded
if file_name is not None:
    # capture image with intended rotation
    image = check_image_rotation(file_name)
    
    # display the image directly with adjusted width
    st.image(image, width=300)  # Adjust width as needed
    
    processing_message = st.empty()
    processing_message.text("Analyzing your image... Please wait.")
    
    # process uploaded image
    emotion_predictions = process_photo(file_name) 
    # process emotion predictions
    result = process_emotions(st.session_state.model, st.session_state.tokenizer, emotion_predictions)
    
    processing_message.empty()
    
    # generate affective journaling prompt based on emotion predictions
    prompt = generate_prompt(result)
    # display journal prompt
    st.write(prompt)