import torch
from transformers import AutoTokenizer, AutoModelForSequenceClassification
import pandas as pd

class EmotionAnalyzer:
    def __init__(self, hub_path="vinpalaka/optimized-go-emotions"):
        """
        Initialize the emotion analyzer with a model from the Hugging Face Hub.
        Args:
            hub_path (str): Path to the model on the Hugging Face Hub
        """
        # Set optimal CPU threads
        torch.set_num_threads(4)
        
        # Load tokenizer and model from Hub
        self.tokenizer = AutoTokenizer.from_pretrained(hub_path, use_fast=True)
        self.model = AutoModelForSequenceClassification.from_pretrained(
            hub_path,
            low_cpu_mem_usage=True
        )
        self.model.eval()
    
    def analyze(self, text, min_score=0.01):
        """
        Analyze text and return all emotions with their scores.
        Args:
            text (str): Text to analyze
            min_score (float): Minimum score to include in results (0-1)
        Returns:
            pandas.DataFrame: Sorted emotions and their scores
        """
        # Tokenize text
        inputs = self.tokenizer(text, return_tensors="pt", padding=True, truncation=True)
        
        # Get predictions
        with torch.no_grad():
            outputs = self.model(**inputs)
            scores = torch.nn.functional.softmax(outputs.logits, dim=-1)[0]
            
        # Convert to DataFrame
        results = []
        for idx, score in enumerate(scores):
            score = score.item()
            if score >= min_score:
                results.append({
                    'emotion': self.model.config.id2label[idx],
                    'score': round(score * 100, 2)  # Convert to percentage
                })
        
        # Create DataFrame and sort by score
        df = pd.DataFrame(results)
        if not df.empty:
            df = df.sort_values('score', ascending=False).reset_index(drop=True)
        
        return df

if __name__ == "__main__":
    # Example usage
    analyzer = EmotionAnalyzer()
    
    # Test text
    text = """Do we just try our best and still be vulnerable to bad news."""
    
    # Get emotions
    results = analyzer.analyze(text)
    
    # Print results
    print(f"\nAnalyzing text: {text}\n")
    print("Emotions detected:")
    print(results.to_string(index=False))