visualizations.py

import plotly.graph_objects as go
import plotly.express as px
from wordcloud import WordCloud
import matplotlib.pyplot as plt
import io
import base64
from typing import List, Dict
import numpy as np
import logging

logger = logging.getLogger(__name__)


class VisualizationGenerator:
    """
    Generate various visualizations for emotion analysis
    """
    
    def __init__(self):
        """Initialize visualization generator"""
        self.emotion_colors = {
            'joy': '#FFD700',
            'sadness': '#4682B4',
            'anger': '#DC143C',
            'fear': '#483D8B',
            'surprise': '#00CED1',
            'neutral': '#A9A9A9',
            'disgust': '#556B2F'
        }
    
    def create_visualizations(self, results: List[Dict]) -> Dict:
        """
        Create all visualizations from emotion analysis results
        
        Args:
            results: List of emotion detection results
            
        Returns:
            Dictionary containing all visualization data
        """
        try:
            # Aggregate emotions
            aggregated = self._aggregate_results(results)
            
            # Create various charts
            emotion_pie = self._create_emotion_pie(aggregated['emotions'])
            emotion_bar = self._create_emotion_bar(aggregated['emotions'])
            sentiment_gauge = self._create_sentiment_gauge(aggregated['sentiment_score'])
            timeline = self._create_emotion_timeline(results)
            radar_chart = self._create_emotion_radar(aggregated['emotions'])
            
            return {
                "emotion_pie": emotion_pie,
                "emotion_bar": emotion_bar,
                "sentiment_gauge": sentiment_gauge,
                "emotion_timeline": timeline,
                "emotion_radar": radar_chart,
                "statistics": {
                    "total_texts": len(results),
                    "dominant_emotion": aggregated['dominant_emotion'],
                    "avg_sentiment": aggregated['sentiment_score'],
                    "sentiment_label": aggregated['sentiment_label']
                }
            }
            
        except Exception as e:
            logger.error(f"Error creating visualizations: {str(e)}")
            return {}
    
    def _aggregate_results(self, results: List[Dict]) -> Dict:
        """Aggregate multiple emotion results"""
        if not results:
            return {}
        
        # Initialize
        emotion_sums = {}
        sentiment_sum = 0.0
        
        # Aggregate
        for result in results:
            for emotion, score in result['emotions'].items():
                emotion_sums[emotion] = emotion_sums.get(emotion, 0) + score
            sentiment_sum += result.get('sentiment_score', 0)
        
        # Calculate averages
        n = len(results)
        emotions_avg = {emotion: score / n for emotion, score in emotion_sums.items()}
        
        # Get dominant
        dominant = max(emotions_avg.items(), key=lambda x: x[1])
        
        # Sentiment
        avg_sentiment = sentiment_sum / n
        sentiment_label = "positive" if avg_sentiment > 0.1 else "negative" if avg_sentiment < -0.1 else "neutral"
        
        return {
            "emotions": emotions_avg,
            "dominant_emotion": dominant[0],
            "sentiment_score": avg_sentiment,
            "sentiment_label": sentiment_label
        }
    
    def _create_emotion_pie(self, emotions: Dict[str, float]) -> str:
        """Create a pie chart of emotion distribution"""
        try:
            # Filter significant emotions
            significant = {k: v for k, v in emotions.items() if v > 0.05}
            
            if not significant:
                significant = emotions
            
            # Create pie chart
            fig = go.Figure(data=[go.Pie(
                labels=[e.capitalize() for e in significant.keys()],
                values=list(significant.values()),
                marker=dict(colors=[self.emotion_colors.get(e, '#808080') for e in significant.keys()]),
                hole=0.3,
                textinfo='label+percent',
                textfont=dict(size=14),
                hovertemplate='<b>%{label}</b><br>%{percent}<extra></extra>'
            )])
            
            fig.update_layout(
                title={
                    'text': 'Emotion Distribution',
                    'x': 0.5,
                    'xanchor': 'center',
                    'font': {'size': 20, 'color': '#333'}
                },
                showlegend=True,
                height=400,
                paper_bgcolor='rgba(0,0,0,0)',
                plot_bgcolor='rgba(0,0,0,0)',
                font=dict(family="Arial, sans-serif")
            )
            
            return fig.to_json()
            
        except Exception as e:
            logger.error(f"Error creating pie chart: {str(e)}")
            return "{}"
    
    def _create_emotion_bar(self, emotions: Dict[str, float]) -> str:
        """Create a bar chart of emotions"""
        try:
            # Sort emotions
            sorted_emotions = sorted(emotions.items(), key=lambda x: x[1], reverse=True)
            
            labels = [e[0].capitalize() for e in sorted_emotions]
            values = [e[1] * 100 for e in sorted_emotions]
            colors = [self.emotion_colors.get(e[0], '#808080') for e in sorted_emotions]
            
            fig = go.Figure(data=[go.Bar(
                x=labels,
                y=values,
                marker=dict(color=colors),
                text=[f'{v:.1f}%' for v in values],
                textposition='auto',
                hovertemplate='<b>%{x}</b><br>%{y:.1f}%<extra></extra>'
            )])
            
            fig.update_layout(
                title={
                    'text': 'Emotion Intensity',
                    'x': 0.5,
                    'xanchor': 'center',
                    'font': {'size': 20, 'color': '#333'}
                },
                xaxis_title='Emotion',
                yaxis_title='Intensity (%)',
                height=400,
                paper_bgcolor='rgba(0,0,0,0)',
                plot_bgcolor='rgba(0,0,0,0)',
                font=dict(family="Arial, sans-serif"),
                yaxis=dict(gridcolor='#E5E5E5')
            )
            
            return fig.to_json()
            
        except Exception as e:
            logger.error(f"Error creating bar chart: {str(e)}")
            return "{}"
    
    def _create_sentiment_gauge(self, sentiment_score: float) -> str:
        """Create a gauge chart for sentiment"""
        try:
            # Normalize to 0-100 scale
            gauge_value = (sentiment_score + 1) * 50
            
            # Determine color
            if sentiment_score > 0.1:
                color = '#28a745'
            elif sentiment_score < -0.1:
                color = '#dc3545'
            else:
                color = '#ffc107'
            
            fig = go.Figure(go.Indicator(
                mode="gauge+number+delta",
                value=gauge_value,
                domain={'x': [0, 1], 'y': [0, 1]},
                title={'text': "Sentiment Score", 'font': {'size': 20}},
                gauge={
                    'axis': {'range': [0, 100], 'tickwidth': 1},
                    'bar': {'color': color},
                    'steps': [
                        {'range': [0, 33], 'color': '#ffebee'},
                        {'range': [33, 66], 'color': '#fff9e6'},
                        {'range': [66, 100], 'color': '#e8f5e9'}
                    ],
                    'threshold': {
                        'line': {'color': "red", 'width': 4},
                        'thickness': 0.75,
                        'value': 50
                    }
                }
            ))
            
            fig.update_layout(
                height=300,
                paper_bgcolor='rgba(0,0,0,0)',
                font=dict(family="Arial, sans-serif")
            )
            
            return fig.to_json()
            
        except Exception as e:
            logger.error(f"Error creating gauge: {str(e)}")
            return "{}"
    
    def _create_emotion_timeline(self, results: List[Dict]) -> str:
        """Create a timeline of emotion changes"""
        try:
            if len(results) < 2:
                return "{}"
            
            # Extract emotion data over time
            emotions_over_time = {}
            for emotion in results[0]['emotions'].keys():
                emotions_over_time[emotion] = [r['emotions'][emotion] * 100 for r in results]
            
            # Create line chart
            fig = go.Figure()
            
            for emotion, values in emotions_over_time.items():
                if max(values) > 5:  # Only show significant emotions
                    fig.add_trace(go.Scatter(
                        x=list(range(1, len(values) + 1)),
                        y=values,
                        mode='lines+markers',
                        name=emotion.capitalize(),
                        line=dict(color=self.emotion_colors.get(emotion, '#808080'), width=2),
                        marker=dict(size=6)
                    ))
            
            fig.update_layout(
                title={
                    'text': 'Emotion Timeline',
                    'x': 0.5,
                    'xanchor': 'center',
                    'font': {'size': 20, 'color': '#333'}
                },
                xaxis_title='Text Entry',
                yaxis_title='Emotion Intensity (%)',
                height=400,
                paper_bgcolor='rgba(0,0,0,0)',
                plot_bgcolor='rgba(0,0,0,0)',
                font=dict(family="Arial, sans-serif"),
                hovermode='x unified',
                yaxis=dict(gridcolor='#E5E5E5'),
                xaxis=dict(gridcolor='#E5E5E5')
            )
            
            return fig.to_json()
            
        except Exception as e:
            logger.error(f"Error creating timeline: {str(e)}")
            return "{}"
    
    def _create_emotion_radar(self, emotions: Dict[str, float]) -> str:
        """Create a radar chart of emotions"""
        try:
            categories = [e.capitalize() for e in emotions.keys()]
            values = [v * 100 for v in emotions.values()]
            
            fig = go.Figure()
            
            fig.add_trace(go.Scatterpolar(
                r=values,
                theta=categories,
                fill='toself',
                fillcolor='rgba(75, 192, 192, 0.2)',
                line=dict(color='rgb(75, 192, 192)', width=2),
                marker=dict(size=8)
            ))
            
            fig.update_layout(
                polar=dict(
                    radialaxis=dict(
                        visible=True,
                        range=[0, max(values) * 1.2] if max(values) > 0 else [0, 100],
                        gridcolor='#E5E5E5'
                    ),
                    angularaxis=dict(gridcolor='#E5E5E5')
                ),
                title={
                    'text': 'Emotion Profile',
                    'x': 0.5,
                    'xanchor': 'center',
                    'font': {'size': 20, 'color': '#333'}
                },
                height=400,
                paper_bgcolor='rgba(0,0,0,0)',
                font=dict(family="Arial, sans-serif"),
                showlegend=False
            )
            
            return fig.to_json()
            
        except Exception as e:
            logger.error(f"Error creating radar chart: {str(e)}")
            return "{}"