import torch
import gradio as gr
from PIL import Image
from transformers import AutoTokenizer, AutoModelForSequenceClassification

js = """
function createGradioAnimation() {
    var container = document.createElement('div');
    container.id = 'gradio-animation';
    container.style.fontSize = '2em';
    container.style.fontWeight = 'bold';
    container.style.textAlign = 'center';
    container.style.marginBottom = '20px';

    var text = 'Stress Prediction Model';
    for (var i = 0; i < text.length; i++) {
        (function(i){
            setTimeout(function(){
                var letter = document.createElement('span');
                letter.style.opacity = '0';
                letter.style.transition = 'opacity 0.5s';
                letter.innerText = text[i];

                container.appendChild(letter);

                setTimeout(function() {
                    letter.style.opacity = '1';
                }, 50);
            }, i * 250);
        })(i);
    }

    var gradioContainer = document.querySelector('.gradio-container');
    gradioContainer.insertBefore(container, gradioContainer.firstChild);

    return 'Animation created';
}
"""

saved_directory = 'jnyx74/stress-prediction'
tokenizer = AutoTokenizer.from_pretrained(saved_directory)
model = AutoModelForSequenceClassification.from_pretrained(saved_directory)

# gr.load("models/jnyx74/stress-prediction").launch()

#"LABEL_0": "I think you don't feel stress. Perhaps, describe more, so I can understand you more!",
#"LABEL_1": "Darling, I sensed that you were stressed. Are you alright?""

background = Image.open('quote.jpg')
with gr.Blocks(js=js,theme=gr.themes.Soft()) as demo:
    gr.Image(background, height = '400px',interactive = False)
    gr.Markdown(
        '''
        # Let me study you, perhaps?
        Not everyone tends to express/ knowing to relieve their stress in a proper way. Therefore, are you stress? Perhaps,
        I could have a guess here. I am a fine-tuned DeepLearning/Transformers Model on DistilBert Model with training on reddit datasets.
        '''
    )
    gr.Markdown("Start typing below and then click **Study Me** to see the output.")
    inp = gr.Text(placeholder = "How do you feel today?", label="Sentence Me:")
    btn = gr.Button("Study Me")
    with gr.Column(visible=False) as output_col:
        out_label = gr.Markdown("# Ooh, I think ...")
        out = gr.Text(label="Result",interactive = False)
    examples = gr.Examples(examples=["By serendipity, I meet her once again and for real, I miss her.", 
                                     "Insomnia and overthinking is really killing me as my final year project is reaching.",
                                     "I just won a lottery and wanting to own a house in Kuching.",
                                     "I can't believe I just hit a car just now and the car driver just ran away, how ridiculous?",
                                     "I hate kids but my wife insists to have one, can't we just adopt?"]
                                    ,
                                      inputs = [inp])
    
    def form_submit(inp):
        if len(inp)<=15:
            gr.Warning("Describe/ Express more, a sentence with expression is more appreciated")
            return {
                output_col: gr.Column(visible=True),
                out: gr.Text(value="Please express more of you!!")}
        else:
            inputs = tokenizer([inp], return_tensors="pt")
            with torch.no_grad():
                logits = model(**inputs).logits

            predicted_class_id = logits.argmax().item()
            result = model.config.id2label[predicted_class_id]
            if result == 'LABEL_0':
                result_value = "I think you don't feel stress. Perhaps, describe more, so I can understand you more!"
            else: result_value = "Darling, I sensed that you were stressed. Are you alright?"
            gr.Info("Success Executed")
            return {
                output_col: gr.Column(visible=True),
                out: gr.Text(value=result_value)
            }
    btn.click(fn=form_submit, inputs=inp, outputs=[out, output_col])
    

demo.launch()