import streamlit as st
import pandas as pd
import numpy as np
import tensorflow as tf
import tensorflow_hub as tf_hub

number_to_feeling = {
  '0': 'sadness',
  '1': 'anger',
  '2': 'love',
  '3': 'surprise',
  '4': 'fear',
  '5': 'joy'
}

def get_feeling(number):
  feeling = number_to_feeling.get(str(number), "Unknown feeling")
  return feeling

# Load the model function
def load_model():
  return tf.keras.models.load_model('model.keras', custom_objects={'KerasLayer': tf_hub.KerasLayer})

def app():
  st.header('Prediction', divider='rainbow')

  user_input = st.text_input("Enter your text here:")
  if 'the_model' not in st.session_state:
    with st.spinner("Loading the model, please wait..."):
      st.session_state.the_model = load_model()
  if st.button('Predict', type="secondary"):
    data = {
      "text_processed": [
          user_input
      ]
    }
    df = pd.DataFrame(data)
    with st.spinner("Making prediction..."):
      # Replace with your preprocessing and prediction code
      predictions = st.session_state.the_model.predict(df)
      predicted_class = np.argmax(predictions, axis=1)
      the_sentiment = predicted_class[0]

      st.write(f"We have predicted that the sentiment of this text is {get_feeling(the_sentiment)}")
  else:
    st.write("Click the button to predict!")