visualizations.py

"""
Visualization components for sentiment analysis results
"""

import plotly.graph_objects as go
import plotly.express as px
from wordcloud import WordCloud
import matplotlib.pyplot as plt
from typing import List, Dict
import pandas as pd
import io


def create_sentiment_pie_chart(stats: Dict) -> go.Figure:
    """
    Create pie chart showing sentiment distribution
    """
    labels = ['Positive', 'Negative', 'Neutral']
    values = [
        stats['positive_count'],
        stats['negative_count'],
        stats['neutral_count']
    ]
    colors = ['#10b981', '#ef4444', '#6b7280']

    fig = go.Figure(data=[go.Pie(
        labels=labels,
        values=values,
        marker=dict(colors=colors),
        hole=0.4,
        textinfo='label+percent',
        textfont_size=14
    )])

    fig.update_layout(
        title={
            'text': 'Sentiment Distribution',
            'x': 0.5,
            'xanchor': 'center',
            'font': {'size': 20, 'color': '#1f2937'}
        },
        showlegend=True,
        height=400,
        margin=dict(t=80, b=20, l=20, r=20)
    )

    return fig


def create_sentiment_bar_chart(results: List[Dict]) -> go.Figure:
    """
    Create bar chart showing sentiment scores for each text
    """
    df = pd.DataFrame(results)

    # Limit to first 20 results for readability
    if len(df) > 20:
        df = df.head(20)

    # Create short labels for x-axis
    df['short_text'] = df['original_text'].str[:30] + '...'

    # Color based on sentiment
    colors = df['sentiment'].map({
        'Positive': '#10b981',
        'Negative': '#ef4444',
        'Neutral': '#6b7280'
    })

    fig = go.Figure(data=[
        go.Bar(
            x=df['short_text'],
            y=df['compound_score'],
            marker_color=colors,
            text=df['compound_score'].round(3),
            textposition='outside',
            hovertemplate='<b>Text:</b> %{customdata}<br>' +
                          '<b>Score:</b> %{y:.4f}<br>' +
                          '<extra></extra>',
            customdata=df['original_text']
        )
    ])

    fig.update_layout(
        title={
            'text': 'Sentiment Scores by Text',
            'x': 0.5,
            'xanchor': 'center',
            'font': {'size': 20, 'color': '#1f2937'}
        },
        xaxis_title='Text',
        yaxis_title='Compound Score',
        height=500,
        margin=dict(t=80, b=120, l=60, r=20),
        xaxis_tickangle=-45,
        yaxis_range=[-1, 1],
        showlegend=False
    )

    # Add horizontal line at y=0
    fig.add_hline(y=0, line_dash="dash", line_color="gray", opacity=0.5)

    return fig


def create_score_distribution_chart(results: List[Dict]) -> go.Figure:
    """
    Create histogram showing distribution of compound scores
    """
    scores = [r['compound_score'] for r in results]

    fig = go.Figure(data=[go.Histogram(
        x=scores,
        nbinsx=30,
        marker_color='#3b82f6',
        opacity=0.7
    )])

    fig.update_layout(
        title={
            'text': 'Distribution of Sentiment Scores',
            'x': 0.5,
            'xanchor': 'center',
            'font': {'size': 20, 'color': '#1f2937'}
        },
        xaxis_title='Compound Score',
        yaxis_title='Frequency',
        height=400,
        margin=dict(t=80, b=60, l=60, r=20),
        showlegend=False
    )

    return fig


def create_wordcloud(texts: List[str], sentiment: str = 'all') -> plt.Figure:
    """
    Create word cloud from texts
    """
    # Combine all texts
    combined_text = ' '.join(texts)

    # Set color based on sentiment
    if sentiment == 'Positive':
        colormap = 'Greens'
    elif sentiment == 'Negative':
        colormap = 'Reds'
    else:
        colormap = 'Blues'

    # Generate word cloud
    wordcloud = WordCloud(
        width=800,
        height=400,
        background_color='white',
        colormap=colormap,
        max_words=100,
        relative_scaling=0.5,
        min_font_size=10
    ).generate(combined_text)

    # Create matplotlib figure
    fig, ax = plt.subplots(figsize=(10, 5))
    ax.imshow(wordcloud, interpolation='bilinear')
    ax.axis('off')
    ax.set_title(f'Word Cloud - {sentiment} Sentiment',
                 fontsize=16, fontweight='bold', pad=20)

    plt.tight_layout()
    return fig


def create_confidence_distribution(results: List[Dict]) -> go.Figure:
    """
    Create pie chart showing confidence level distribution
    """
    confidence_counts = {}
    for result in results:
        conf = result['confidence']
        confidence_counts[conf] = confidence_counts.get(conf, 0) + 1

    labels = list(confidence_counts.keys())
    values = list(confidence_counts.values())
    colors = {'High': '#10b981', 'Medium': '#f59e0b', 'Low': '#ef4444'}
    color_list = [colors.get(label, '#6b7280') for label in labels]

    fig = go.Figure(data=[go.Pie(
        labels=labels,
        values=values,
        marker=dict(colors=color_list),
        textinfo='label+percent',
        textfont_size=14
    )])

    fig.update_layout(
        title={
            'text': 'Confidence Level Distribution',
            'x': 0.5,
            'xanchor': 'center',
            'font': {'size': 20, 'color': '#1f2937'}
        },
        showlegend=True,
        height=400,
        margin=dict(t=80, b=20, l=20, r=20)
    )

    return fig