import pandas as pd
import numpy
# import tensorflow as tf
from tensorflow.keras.layers import TFSMLayer
from transformers import TFAutoModelForSequenceClassification, AutoTokenizer
import streamlit as st


# model = tf.keras.models.load_model(f'https://huggingface.co/wismaeka/itsok/resolve/main/')

# model = TFSMLayer('/itsok', call_endpoint='serving_default')
model_name='wismaeka/itsok'

model = TFAutoModelForSequenceClassification.from_pretrained(model_name)
tokenizer = AutoTokenizer.from_pretrained(model_name)


def run():
    st.title('How are you feeling today?')
    st.write('This is a simple web app to predict sentiment of a text using deep learning. Input your feeling below to get the prediction.')
    st.write('Trust me, I have analyzed it for you!')

    texts = st.text_input('Text', 'I feel so sad today')

    def convert_to_label(pred):
        if pred == 0:
            return 'Normal'
        elif pred == 1:
            return 'Suicidal'
        elif pred == 2:
            return 'Anxiety'
        elif pred == 3:
            return 'Depression'
        elif pred == 4:
            return 'Stress'
        elif pred == 5:
            return 'Bipolar'
        elif pred == 6:
            return 'Personality Disorder'
        else:
            return 'Unknown'

    if st.button("Predict Your Feeling"):
        # prediction = model.predict(text)
        inputs = tokenizer(texts, return_tensors="tf", padding=True, truncation=True)
        
        #Perform inference
        outputs = model(inputs)
        logits = outputs.logits
        
        # If you want to get the predicted classes
        prediction = tf.argmax(logits, axis=-1).numpy()

        label = convert_to_label(prediction)
        if label == 'Normal':
            st.success('Hi! Keep up the good work! You are feeling Okay today.')
        elif label == 'Suicidal':
            st.error('Hi! I detect you are feeling Suicidal. Please seek help.')
        elif label == 'Anxiety':
            st.error('Hi! I detect you are feeling Anxious. You may want to talk to someone.')
        elif label == 'Depression':
            st.error('Hi! I detect you are feeling Depressed. Please seek help.')
        elif label == 'Stress':
            st.error('Hi! I detect you are feeling Stressed. Please take a break.')
        elif label == 'Bipolar':
            st.error('Hi! I detect you are feeling Bipolar. Please seek help.')
        elif label == 'Personality Disorder':
            st.error('Hi! I detect you are having Personality Disorder. Please seek help.')
        else:
            st.error('Hi! I cannot detect your feeling. Please try again.')


if __name__ == '__main__':
    run()