backend/src/visualize.py

"""Visualization module for emotion analysis results."""

import matplotlib.pyplot as plt
import seaborn as sns
import numpy as np
from pathlib import Path
from typing import Dict, Optional
from datetime import datetime


class EmotionVisualizer:
    """Generate visualizations for emotion analysis results."""

    def __init__(self, output_dir: Path = None):
        """
        Initialize the visualizer.

        Args:
            output_dir: Directory to save visualization files
        """
        self.output_dir = output_dir or Path("visualizations")
        self.output_dir.mkdir(exist_ok=True)

        # Set style
        plt.style.use('default')
        sns.set_palette("husl")

        # Emotion colors for consistency
        self.emotion_colors = {
            'joy': '#FFD700',      # Gold
            'sadness': '#4169E1',  # Royal Blue
            'anger': '#DC143C',    # Crimson
            'fear': '#9370DB',     # Medium Purple
            'love': '#FF69B4',     # Hot Pink
            'surprise': '#FF8C00', # Dark Orange
        }

    def create_emotion_bar_chart(self,
                                 probabilities: Dict[str, float],
                                 text: str = "",
                                 save_path: Optional[Path] = None,
                                 show_chart: bool = True,
                                 primary_emotion: Optional[str] = None) -> Path:
        """
        Create a bar chart showing emotion probabilities.

        Args:
            probabilities: Dict of emotion -> probability
            text: Input text (for title)
            save_path: Where to save the chart
            show_chart: Whether to display the chart
            primary_emotion: Explicit primary emotion to override argmax

        Returns:
            Path to saved chart image
        """
        # Prepare data
        emotions = list(probabilities.keys())
        probs = list(probabilities.values())
        colors = [self.emotion_colors.get(emo, '#808080') for emo in emotions]

        # Create figure
        plt.figure(figsize=(12, 8))

        # Create bars
        bars = plt.bar(emotions, probs, color=colors, alpha=0.8, edgecolor='black', linewidth=1)

        # Customize chart
        plt.title(f'Emotion Analysis Results\n"{text[:60]}{"..." if len(text) > 60 else ""}"',
                 fontsize=16, fontweight='bold', pad=20)
        plt.xlabel('Emotions', fontsize=14, fontweight='bold')
        plt.ylabel('Probability', fontsize=14, fontweight='bold')

        # Add percentage labels on bars
        for bar, prob in zip(bars, probs):
            height = bar.get_height()
            plt.text(bar.get_x() + bar.get_width()/2., height + 0.01,
                    f'{prob:.1%}', ha='center', va='bottom',
                    fontweight='bold', fontsize=12)

        # Customize appearance
        plt.ylim(0, max(probs) * 1.2)
        plt.xticks(rotation=45, fontsize=12)
        plt.yticks(fontsize=12)
        plt.grid(axis='y', alpha=0.3, linestyle='--')

        # Add emotion indicators
        if primary_emotion is None:
            max_emotion = max(probabilities, key=probabilities.get)
        else:
            max_emotion = primary_emotion
            
        max_prob = probabilities[max_emotion]
        plt.text(0.02, 0.98, f'Primary: {max_emotion.upper()} ({max_prob:.1%})',
                transform=plt.gca().transAxes, fontsize=14, fontweight='bold',
                bbox=dict(boxstyle="round,pad=0.3", facecolor="lightblue", alpha=0.7),
                verticalalignment='top')

        plt.tight_layout()

        # Save chart
        if save_path is None:
            timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
            save_path = self.output_dir / f"emotion_analysis_{timestamp}.png"

        plt.savefig(save_path, dpi=300, bbox_inches='tight',
                   facecolor='white', edgecolor='none')

        if show_chart:
            plt.show()
        else:
            plt.close()

        return save_path

    def create_comparison_chart(self,
                               results_list: list,
                               titles: list = None,
                               save_path: Optional[Path] = None,
                               show_chart: bool = True) -> Path:
        """
        Create a comparison chart for multiple predictions.

        Args:
            results_list: List of prediction results
            titles: List of titles for each prediction
            save_path: Where to save the chart
            show_chart: Whether to display the chart

        Returns:
            Path to saved chart image
        """
        if not results_list:
            raise ValueError("No results provided for comparison")

        # Prepare data
        emotions = list(results_list[0]['probabilities'].keys())
        n_samples = len(results_list)

        if titles is None:
            titles = [f"Sample {i+1}" for i in range(n_samples)]

        # Create figure with subplots
        fig, axes = plt.subplots(1, n_samples, figsize=(6*n_samples, 6))
        if n_samples == 1:
            axes = [axes]

        fig.suptitle('Emotion Analysis Comparison', fontsize=18, fontweight='bold')

        for idx, (result, title, ax) in enumerate(zip(results_list, titles, axes)):
            probs = list(result['probabilities'].values())
            colors = [self.emotion_colors.get(emo, '#808080') for emo in emotions]

            bars = ax.bar(emotions, probs, color=colors, alpha=0.8,
                         edgecolor='black', linewidth=1)

            # Add labels
            for bar, prob in zip(bars, probs):
                height = bar.get_height()
                ax.text(bar.get_x() + bar.get_width()/2., height + 0.01,
                       f'{prob:.1%}', ha='center', va='bottom',
                       fontweight='bold', fontsize=10)

            ax.set_title(f'{title}\n🎯 {result["emotion"].upper()} ({result["confidence"]:.1%})')
            ax.set_ylim(0, 1)
            ax.tick_params(axis='x', rotation=45)
            ax.grid(axis='y', alpha=0.3, linestyle='--')

        plt.tight_layout()

        # Save chart
        if save_path is None:
            timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
            save_path = self.output_dir / f"emotion_comparison_{timestamp}.png"

        plt.savefig(save_path, dpi=300, bbox_inches='tight',
                   facecolor='white', edgecolor='none')

        if show_chart:
            plt.show()
        else:
            plt.close()

        return save_path

    def create_detailed_analysis_chart(self,
                                      result: dict,
                                      text: str = "",
                                      save_path: Optional[Path] = None,
                                      show_chart: bool = True) -> Path:
        """
        Create a simplified detailed analysis chart with only the bar chart.

        Args:
            result: Prediction result dictionary
            text: Input text
            save_path: Where to save the chart
            show_chart: Whether to display the chart

        Returns:
            Path to saved chart image
        """
        fig = plt.figure(figsize=(14, 8))

        # Single bar chart with enhanced styling
        emotions = list(result['probabilities'].keys())
        probs = list(result['probabilities'].values())
        colors = [self.emotion_colors.get(emo, '#808080') for emo in emotions]

        # Verify probabilities sum to 100%
        total_prob = sum(probs)
        if abs(total_prob - 1.0) > 0.001:
            print(f"⚠️ Warning: Probabilities sum to {total_prob:.4f}, not 1.0")

        bars = plt.bar(emotions, probs, color=colors, alpha=0.8,
                      edgecolor='black', linewidth=2, width=0.6)

        # Add percentage labels on bars
        for bar, prob in zip(bars, probs):
            height = bar.get_height()
            plt.text(bar.get_x() + bar.get_width()/2., height + 0.01,
                    f'{prob:.1%}', ha='center', va='bottom',
                    fontweight='bold', fontsize=14)

        # Enhanced styling
        plt.title(f'Emotion Analysis Results\n"{text[:80]}{"..." if len(text) > 80 else ""}"',
                 fontsize=18, fontweight='bold', pad=20)
        plt.xlabel('Emotions', fontsize=16, fontweight='bold')
        plt.ylabel('Probability', fontsize=16, fontweight='bold')

        # Set y-axis to show 0-100%
        plt.ylim(0, 1.0)
        plt.gca().yaxis.set_major_formatter(plt.FuncFormatter(lambda y, _: f'{y:.0%}'))

        # Customize appearance
        plt.xticks(rotation=0, fontsize=13, fontweight='bold')
        plt.yticks(fontsize=12)
        plt.grid(axis='y', alpha=0.3, linestyle='--')

        # Add primary emotion indicator
        max_emotion = max(result['probabilities'], key=result['probabilities'].get)
        max_prob = result['probabilities'][max_emotion]
        plt.text(0.02, 0.98, f'Primary Emotion: {max_emotion.upper()} ({max_prob:.1%})',
                transform=plt.gca().transAxes, fontsize=16, fontweight='bold',
                bbox=dict(boxstyle="round,pad=0.5", facecolor="lightblue", alpha=0.8),
                verticalalignment='top')

        # Add verification text
        plt.text(0.98, 0.02, f'Total: {sum(probs):.1%}',
                transform=plt.gca().transAxes, fontsize=12,
                bbox=dict(boxstyle="round,pad=0.3", facecolor="lightgreen", alpha=0.7),
                horizontalalignment='right', verticalalignment='bottom')

        plt.tight_layout()

        # Save chart
        if save_path is None:
            timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
            save_path = self.output_dir / f"detailed_analysis_{timestamp}.png"

        plt.savefig(save_path, dpi=300, bbox_inches='tight',
                   facecolor='white', edgecolor='none')

        if show_chart:
            plt.show()
        else:
            plt.close()

        return save_path


def create_quick_chart(probabilities: dict, text: str = "", show: bool = True) -> Path:
    """Quick function to create a simple emotion bar chart."""
    visualizer = EmotionVisualizer()
    return visualizer.create_emotion_bar_chart(probabilities, text, show_chart=show)


if __name__ == "__main__":
    # Demo visualization
    sample_probs = {
        'joy': 0.7,
        'sadness': 0.1,
        'anger': 0.08,
        'fear': 0.06,
        'love': 0.04,
        'surprise': 0.02
    }

    visualizer = EmotionVisualizer()
    chart_path = visualizer.create_emotion_bar_chart(
        sample_probs,
        "I'm feeling so happy today, everything is wonderful!"
    )
    print(f"📊 Chart saved to: {chart_path}")