import matplotlib.pyplot as plt
import numpy as np
from wordcloud import WordCloud
from collections import Counter
from typing import List, Dict, Tuple, Optional
import gc

from config import config
from utils import handle_errors, managed_figure


class ThemeContext:
    """Theme management context"""
    def __init__(self, theme: str = 'default'):
        self.theme = theme
        self.colors = config.THEMES.get(theme, config.THEMES['default'])


class PlotFactory:
    """Factory for creating plots with proper memory management"""
    
    @staticmethod
    @handle_errors(default_return=None)
    def create_sentiment_bars(probs: np.ndarray, theme: ThemeContext) -> plt.Figure:
        """Create sentiment probability bars"""
        with managed_figure(figsize=config.FIGURE_SIZE_SINGLE) as fig:
            ax = fig.add_subplot(111)
            labels = ["Negative", "Positive"]
            colors = [theme.colors['neg'], theme.colors['pos']]
            
            bars = ax.bar(labels, probs, color=colors, alpha=0.8)
            ax.set_title("Sentiment Probabilities", fontweight='bold')
            ax.set_ylabel("Probability")
            ax.set_ylim(0, 1)
            
            for bar, prob in zip(bars, probs):
                ax.text(bar.get_x() + bar.get_width()/2., bar.get_height() + 0.02,
                       f'{prob:.3f}', ha='center', va='bottom', fontweight='bold')
            
            fig.tight_layout()
            return fig
    
    @staticmethod
    @handle_errors(default_return=None)
    def create_confidence_gauge(confidence: float, sentiment: str, theme: ThemeContext) -> plt.Figure:
        """Create confidence gauge"""
        with managed_figure(figsize=config.FIGURE_SIZE_SINGLE) as fig:
            ax = fig.add_subplot(111)
            
            theta = np.linspace(0, np.pi, 100)
            colors = [theme.colors['neg'] if i < 50 else theme.colors['pos'] for i in range(100)]
            
            for i in range(len(theta)-1):
                ax.fill_between([theta[i], theta[i+1]], [0, 0], [0.8, 0.8], 
                               color=colors[i], alpha=0.7)
            
            pos = np.pi * (0.5 + (0.4 if sentiment == 'Positive' else -0.4) * confidence)
            ax.plot([pos, pos], [0, 0.6], 'k-', linewidth=6)
            ax.plot(pos, 0.6, 'ko', markersize=10)
            
            ax.set_xlim(0, np.pi)
            ax.set_ylim(0, 1)
            ax.set_title(f'{sentiment} - Confidence: {confidence:.3f}', fontweight='bold')
            ax.set_xticks([0, np.pi/2, np.pi])
            ax.set_xticklabels(['Negative', 'Neutral', 'Positive'])
            ax.axis('off')
            
            fig.tight_layout()
            return fig
    
    @staticmethod
    @handle_errors(default_return=None)
    def create_lime_keyword_chart(lime_words: List[Tuple[str, float]], sentiment: str, theme: ThemeContext) -> Optional[plt.Figure]:
        """Create horizontal bar chart for LIME key contributing words"""
        if not lime_words:
            return None
        
        with managed_figure(figsize=config.FIGURE_SIZE_SINGLE) as fig:
            ax = fig.add_subplot(111)
            
            words = [word for word, score in lime_words]
            scores = [score for word, score in lime_words]
            
            color = theme.colors['pos'] if sentiment == 'Positive' else theme.colors['neg']
            
            bars = ax.barh(range(len(words)), scores, color=color, alpha=0.7)
            ax.set_yticks(range(len(words)))
            ax.set_yticklabels(words)
            ax.set_xlabel('LIME Attention Weight')
            ax.set_title(f'LIME: Top Contributing Words ({sentiment})', fontweight='bold')
            
            for i, (bar, score) in enumerate(zip(bars, scores)):
                ax.text(bar.get_width() + 0.001, bar.get_y() + bar.get_height()/2.,
                       f'{score:.3f}', ha='left', va='center', fontsize=9)
            
            ax.invert_yaxis()
            ax.grid(axis='x', alpha=0.3)
            fig.tight_layout()
            return fig
    
    @staticmethod
    @handle_errors(default_return=None)
    def create_shap_keyword_chart(shap_words: List[Tuple[str, float]], sentiment: str, theme: ThemeContext) -> Optional[plt.Figure]:
        """Create horizontal bar chart for SHAP key contributing words"""
        if not shap_words:
            return None
        
        with managed_figure(figsize=config.FIGURE_SIZE_SINGLE) as fig:
            ax = fig.add_subplot(111)
            
            words = [word for word, score in shap_words]
            scores = [score for word, score in shap_words]
            
            color = theme.colors['pos'] if sentiment == 'Positive' else theme.colors['neg']
            
            bars = ax.barh(range(len(words)), scores, color=color, alpha=0.7)
            ax.set_yticks(range(len(words)))
            ax.set_yticklabels(words)
            ax.set_xlabel('SHAP Value')
            ax.set_title(f'SHAP: Top Contributing Words ({sentiment})', fontweight='bold')
            
            for i, (bar, score) in enumerate(zip(bars, scores)):
                ax.text(bar.get_width() + 0.001, bar.get_y() + bar.get_height()/2.,
                       f'{score:.3f}', ha='left', va='center', fontsize=9)
            
            ax.invert_yaxis()
            ax.grid(axis='x', alpha=0.3)
            fig.tight_layout()
            return fig
    
    @staticmethod
    @handle_errors(default_return=None)
    def create_wordcloud(text: str, sentiment: str, theme: ThemeContext) -> Optional[plt.Figure]:
        """Create word cloud"""
        if len(text.split()) < 3:
            return None
        
        colormap = 'Greens' if sentiment == 'Positive' else 'Reds'
        wc = WordCloud(width=800, height=400, background_color='white',
                      colormap=colormap, max_words=30).generate(text)
        
        with managed_figure(figsize=config.WORDCLOUD_SIZE) as fig:
            ax = fig.add_subplot(111)
            ax.imshow(wc, interpolation='bilinear')
            ax.axis('off')
            ax.set_title(f'{sentiment} Word Cloud', fontweight='bold')
            fig.tight_layout()
            return fig
    
    @staticmethod
    @handle_errors(default_return=None)
    def create_batch_analysis(results: List[Dict], theme: ThemeContext) -> plt.Figure:
        """Create comprehensive batch visualization"""
        with managed_figure(figsize=config.FIGURE_SIZE_BATCH) as fig:
            gs = fig.add_gridspec(2, 2, hspace=0.3, wspace=0.3)
            
            # Sentiment distribution
            ax1 = fig.add_subplot(gs[0, 0])
            sent_counts = Counter([r['sentiment'] for r in results])
            colors = [theme.colors['pos'], theme.colors['neg']]
            ax1.pie(sent_counts.values(), labels=sent_counts.keys(), 
                   autopct='%1.1f%%', colors=colors[:len(sent_counts)])
            ax1.set_title('Sentiment Distribution')
            
            # Confidence histogram
            ax2 = fig.add_subplot(gs[0, 1])
            confs = [r['confidence'] for r in results]
            ax2.hist(confs, bins=8, alpha=0.7, color='skyblue', edgecolor='black')
            ax2.set_title('Confidence Distribution')
            ax2.set_xlabel('Confidence')
            
            # Sentiment over time
            ax3 = fig.add_subplot(gs[1, :])
            pos_probs = [r['pos_prob'] for r in results]
            indices = range(len(results))
            colors_scatter = [theme.colors['pos'] if r['sentiment'] == 'Positive' 
                            else theme.colors['neg'] for r in results]
            ax3.scatter(indices, pos_probs, c=colors_scatter, alpha=0.7, s=60)
            ax3.axhline(y=0.5, color='gray', linestyle='--', alpha=0.5)
            ax3.set_title('Sentiment Progression')
            ax3.set_xlabel('Review Index')
            ax3.set_ylabel('Positive Probability')
            
            return fig