"""
Gradio Deployment App for Emotion Classification Model
This app loads the trained model from HuggingFace Hub and creates an interactive interface.
"""

import gradio as gr
import torch
from transformers import AutoTokenizer, AutoModelForSequenceClassification
import numpy as np
from huggingface_hub import hf_hub_download

# Configuration
HF_REPO_ID = "hrshlgunjal/emotion-classifier-deberta-v3"
LABELS = ["anger", "fear", "joy", "sadness", "surprise"]
MAX_LEN = 128

# Load model and tokenizer
print("Loading model from HuggingFace Hub...")
model = AutoModelForSequenceClassification.from_pretrained(
    HF_REPO_ID,
    num_labels=len(LABELS),
    problem_type="multi_label_classification"
)
tokenizer = AutoTokenizer.from_pretrained(HF_REPO_ID)
model.eval()
print("✅ Model loaded successfully!")

# Load optimized thresholds
print("Loading optimized thresholds...")
try:
    threshold_path = hf_hub_download(
        repo_id=HF_REPO_ID,
        filename="best_thresholds.npy"
    )
    thresholds = np.load(threshold_path)
    print("✅ Optimized thresholds loaded!")
except Exception as e:
    print(f"⚠️  Could not load thresholds: {e}")
    print("Using default thresholds of 0.5")
    thresholds = np.array([0.5] * len(LABELS))

# Prediction function
def predict_emotions(text):
    """
    Predict emotions from input text.
    
    Args:
        text (str): Input text to analyze
        
    Returns:
        dict: Dictionary of emotion labels and their probabilities
    """
    if not text or not text.strip():
        return {label: 0.0 for label in LABELS}
    
    # Tokenize
    encoding = tokenizer(
        text,
        max_length=MAX_LEN,
        padding='max_length',
        truncation=True,
        return_tensors='pt'
    )
    
    # Predict
    with torch.no_grad():
        outputs = model(**encoding)
        probabilities = torch.sigmoid(outputs.logits).squeeze().numpy()
    
    # Apply thresholds for binary predictions
    predictions = (probabilities >= thresholds).astype(int)
    
    # Create results dictionary
    results = {}
    for label, prob, pred in zip(LABELS, probabilities, predictions):
        results[label] = float(prob)
    
    return results

# Create Gradio Interface
demo = gr.Interface(
    fn=predict_emotions,
    inputs=gr.Textbox(
        lines=5,
        placeholder="Enter text to analyze emotions...",
        label="Input Text"
    ),
    outputs=gr.Label(
        num_top_classes=5,
        label="Emotion Probabilities"
    ),
    title="🎭 Emotion Classifier",
    description="""
    This model classifies text into 5 emotions: **anger**, **fear**, **joy**, **sadness**, and **surprise**.
    
    It uses a fine-tuned DeBERTa-v3 model trained on emotion classification data.
    Enter any text and the model will predict the emotion probabilities!
    """,
    examples=[
        ["I'm so excited about my vacation next week!"],
        ["This makes me really angry and frustrated."],
        ["I'm worried about the upcoming exam."],
        ["That movie made me cry, it was so touching."],
        ["Wow, I didn't expect that to happen!"]
    ],
    theme=gr.themes.Soft(),
    analytics_enabled=False
)

if __name__ == "__main__":
    demo.launch(share=True)