import streamlit as st
from keras.models import load_model
from tensorflow.keras.preprocessing.text import tokenizer_from_json
import contractions
import re
import nltk
import numpy as np
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")

# Download stopwords if not already downloaded
nltk.download('stopwords')

# 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])
    # Pad sequences (optional, depending on your model's input requirements)
    padded_sequences = np.array(sequences)  # Convert list to NumPy array
    return padded_sequences

class_labels = {
    1: "Anxiety",
    2: "Normal",
    3: "Depression",
    4: "Suicidal",
    5: "Stress",
    6: "Bipolar",
    7: "Personality disorder"
}

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_healths_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)[0] 

        # Ensure the predicted_label is always assigned
        predicted_label = class_labels.get(predicted_class, "Unknown")

        # Display the predicted index and label
        st.write("Predicted Index:", predicted_class)
        st.write("Predicted Mental State:", predicted_label)

if __name__ == "__main__":
    main()