import plotly.graph_objects as go
import plotly.express as px
from plotly.subplots import make_subplots
import numpy as np
from collections import Counter
from typing import List, Dict, Optional

from models import handle_errors, ThemeContext

# Optimized Plotly Visualization System
class PlotlyVisualizer:
    """Enhanced Plotly visualizations"""
    
    @staticmethod
    @handle_errors(default_return=None)
    def create_sentiment_gauge(result: Dict, theme: ThemeContext) -> go.Figure:
        """Create animated sentiment gauge"""
        colors = theme.colors
        
        if result.get('has_neutral', False):
            # Three-way gauge
            fig = go.Figure(go.Indicator(
                mode="gauge+number+delta",
                value=result['pos_prob'] * 100,
                domain={'x': [0, 1], 'y': [0, 1]},
                title={'text': f"Sentiment: {result['sentiment']}"},
                delta={'reference': 50},
                gauge={
                    'axis': {'range': [None, 100]},
                    'bar': {'color': colors['pos'] if result['sentiment'] == 'Positive' else colors['neg']},
                    'steps': [
                        {'range': [0, 33], 'color': colors['neg']},
                        {'range': [33, 67], 'color': colors['neu']},
                        {'range': [67, 100], 'color': colors['pos']}
                    ],
                    'threshold': {
                        'line': {'color': "red", 'width': 4},
                        'thickness': 0.75,
                        'value': 90
                    }
                }
            ))
        else:
            # Two-way gauge
            fig = go.Figure(go.Indicator(
                mode="gauge+number",
                value=result['confidence'] * 100,
                domain={'x': [0, 1], 'y': [0, 1]},
                title={'text': f"Confidence: {result['sentiment']}"},
                gauge={
                    'axis': {'range': [None, 100]},
                    'bar': {'color': colors['pos'] if result['sentiment'] == 'Positive' else colors['neg']},
                    'steps': [
                        {'range': [0, 50], 'color': "lightgray"},
                        {'range': [50, 100], 'color': "gray"}
                    ]
                }
            ))
        
        fig.update_layout(height=400, font={'size': 16})
        return fig
    
    @staticmethod
    @handle_errors(default_return=None)
    def create_probability_bars(result: Dict, theme: ThemeContext) -> go.Figure:
        """Create probability bar chart"""
        colors = theme.colors
        
        if result.get('has_neutral', False):
            labels = ['Negative', 'Neutral', 'Positive']
            values = [result['neg_prob'], result['neu_prob'], result['pos_prob']]
            bar_colors = [colors['neg'], colors['neu'], colors['pos']]
        else:
            labels = ['Negative', 'Positive']
            values = [result['neg_prob'], result['pos_prob']]
            bar_colors = [colors['neg'], colors['pos']]
        
        fig = go.Figure(data=[
            go.Bar(x=labels, y=values, marker_color=bar_colors, 
                   text=[f'{v:.3f}' for v in values], textposition='outside')
        ])
        
        fig.update_layout(
            title="Sentiment Probabilities",
            yaxis_title="Probability",
            height=400,
            showlegend=False
        )
        return fig
    
    @staticmethod
    @handle_errors(default_return=None)
    def create_batch_summary(results: List[Dict], theme: ThemeContext) -> go.Figure:
        """Create batch analysis summary"""
        colors = theme.colors
        
        # Count sentiments
        sentiments = [r['sentiment'] for r in results if 'sentiment' in r and r['sentiment'] != 'Error']
        sentiment_counts = Counter(sentiments)
        
        # Create pie chart
        fig = go.Figure(data=[go.Pie(
            labels=list(sentiment_counts.keys()),
            values=list(sentiment_counts.values()),
            marker_colors=[colors.get(s.lower()[:3], '#999999') for s in sentiment_counts.keys()],
            textinfo='label+percent',
            hole=0.3
        )])
        
        fig.update_layout(
            title=f"Batch Analysis Summary ({len(results)} texts)",
            height=400
        )
        return fig
    
    @staticmethod
    @handle_errors(default_return=None)
    def create_confidence_distribution(results: List[Dict]) -> go.Figure:
        """Create confidence distribution plot"""
        confidences = [r['confidence'] for r in results if 'confidence' in r and r['sentiment'] != 'Error']
        if not confidences:
            return go.Figure()
        
        fig = go.Figure(data=[go.Histogram(
            x=confidences,
            nbinsx=20,
            marker_color='skyblue',
            opacity=0.7
        )])

        fig.update_layout(
            title="Confidence Distribution",
            xaxis_title="Confidence Score",
            yaxis_title="Frequency",
            height=400
        )
        return fig
    
    @staticmethod
    @handle_errors(default_return=None)
    def create_history_dashboard(history: List[Dict], theme: ThemeContext) -> go.Figure:
        """Create comprehensive history dashboard"""
        if len(history) < 2:
            return go.Figure()
        
        # Create subplots
        fig = make_subplots(
            rows=2, cols=2,
            subplot_titles=['Sentiment Timeline', 'Confidence Distribution', 
                           'Language Distribution', 'Sentiment Summary'],
            specs=[[{"secondary_y": False}, {"secondary_y": False}],
                   [{"type": "pie"}, {"type": "bar"}]]
        )
        
        # Extract data
        indices = list(range(len(history)))
        pos_probs = [item.get('pos_prob', 0) for item in history]
        confidences = [item['confidence'] for item in history]
        sentiments = [item['sentiment'] for item in history]
        languages = [item.get('language', 'en') for item in history]
        
        # Sentiment timeline
        colors_map = {'Positive': theme.colors['pos'], 'Negative': theme.colors['neg'], 'Neutral': theme.colors['neu']}
        colors = [colors_map.get(s, '#999999') for s in sentiments]
        
        fig.add_trace(
            go.Scatter(x=indices, y=pos_probs, mode='lines+markers', 
                      marker=dict(color=colors, size=8),
                      name='Positive Probability'),
            row=1, col=1
        )
        
        # Confidence distribution
        fig.add_trace(
            go.Histogram(x=confidences, nbinsx=10, name='Confidence'),
            row=1, col=2
        )
        
        # Language distribution
        lang_counts = Counter(languages)
        fig.add_trace(
            go.Pie(labels=list(lang_counts.keys()), values=list(lang_counts.values()),
                   name="Languages"),
            row=2, col=1
        )
        
        # Sentiment summary
        sent_counts = Counter(sentiments)
        sent_colors = [colors_map.get(k, '#999999') for k in sent_counts.keys()]
        fig.add_trace(
            go.Bar(x=list(sent_counts.keys()), y=list(sent_counts.values()),
                   marker_color=sent_colors),
            row=2, col=2
        )
        
        fig.update_layout(height=800, showlegend=False)
        return fig