import streamlit as st
from keras.models import load_model
from tensorflow.keras.preprocessing.text import tokenizer_from_json
import contractions
import re
from nltk.corpus import stopwords
import json

# Set page configuration
st.set_page_config(page_title="Mental Health Classification")

# Page title
st.title("Mental Health Classification")

# Load the tokenizer (make sure the tokenizer file is in the correct path)
def load_tokenizer():
    with open('tokenizer.json') as f:
        tokenizer_json = json.load(f)
    return tokenizer_from_json(tokenizer_json)

# Preprocess text function
def preprocess_text(input_text, tokenizer):
    text = contractions.fix(input_text)
    text = re.sub(r"[^a-z\s]", "", text)
    text = text.lower()
    # Tokenize the words
    words = text.split()
    # Remove stopwords
    stop_words = set(stopwords.words('english'))
    words = [word for word in words if word not in stop_words]
    clean_text = " ".join(words)
    # Convert to sequences
    sequences = tokenizer.texts_to_sequences([clean_text])
    return sequences

def main():
    # Text input for mental state
    input_text = st.text_input("Enter the Mental state here...")

    # Submit button
    submit_button = st.button("Classify")

    if submit_button and input_text:
        # Load the model and tokenizer
        model = load_model("mental_health_model.h5")
        tokenizer = load_tokenizer()

        # Preprocess the input text
        processed_text = preprocess_text(input_text, tokenizer)

        # Make prediction
        response = model.predict(processed_text)
        predicted_class = response.argmax(axis=-1) 
        # Display the prediction result
        st.write("Predicted Mental State:", predicted_class)

if __name__ == "__main__":
    main()