import os
import torch
from transformers import AutoTokenizer, AutoModelForSequenceClassification
import gradio as gr

os.environ["TOKENIZERS_PARALLELISM"] = "true"

emotion_tokenizer = AutoTokenizer.from_pretrained("j-hartmann/emotion-english-distilroberta-base")
emotion_model = AutoModelForSequenceClassification.from_pretrained("j-hartmann/emotion-english-distilroberta-base")
emotion_labels = ["anger", "disgust", "fear", "joy", "neutral", "sadness", "surprise"]

def analyze_emotion(text):
    try:
        inputs = emotion_tokenizer(text, return_tensors="pt", truncation=True, padding=True, max_length=512)
        outputs = emotion_model(**inputs)
        probs = torch.nn.functional.softmax(outputs.logits, dim=-1)
        max_prob, max_index = torch.max(probs, dim=1)
        return emotion_labels[max_index.item()], f"{max_prob.item():.4f}"
    except Exception as e:
        print(f"Error in emotion analysis: {e}")
        return "Error", "N/A"

def create_emotion_tab():
    with gr.Row():
        with gr.Column(scale=2):
            input_text = gr.Textbox(value='I actually speak to the expets myself to give you the best value you can get', lines=5, placeholder="Enter text here...", label="Input Text")
            with gr.Row():
                clear_btn = gr.Button("Clear", scale=1)
                submit_btn = gr.Button("Analyze", scale=1, elem_classes="submit")
        with gr.Column(scale=1):
            output_emotion = gr.Textbox(label="Detected Emotion")
            output_confidence = gr.Textbox(label="Emotion Confidence Score")
    
    submit_btn.click(analyze_emotion, inputs=[input_text], outputs=[output_emotion, output_confidence])
    clear_btn.click(lambda: ("", "", ""), outputs=[input_text, output_emotion, output_confidence])
    gr.Examples(["I am so happy today!", "I feel terrible and sad.", "This is a neutral statement."], inputs=[input_text])