src/explainability/attention_visualization.py

"""
Attention Visualization
=======================
Visualize attention weights from deep learning models for interpretability.
"""

import numpy as np
from typing import Dict, List, Optional, Tuple, Any
from dataclasses import dataclass
import logging
from pathlib import Path
import json

logger = logging.getLogger(__name__)


@dataclass
class AttentionConfig:
    """Configuration for attention visualization."""
    output_dir: str = "visualizations"
    colormap: str = "viridis"
    figure_size: Tuple[int, int] = (12, 8)
    dpi: int = 100


class AttentionVisualizer:
    """
    Visualizes attention weights from transformer and attention-based models.
    
    Supports:
    - Self-attention heatmaps
    - Feature importance from attention
    - Multi-head attention analysis
    - Temporal attention patterns
    """
    
    def __init__(self, config: AttentionConfig = None):
        self.config = config or AttentionConfig()
        self.output_dir = Path(self.config.output_dir)
        self.output_dir.mkdir(parents=True, exist_ok=True)
        self._has_matplotlib = False
        self._check_dependencies()
        
    def _check_dependencies(self):
        """Check if visualization dependencies are available."""
        try:
            import matplotlib
            matplotlib.use('Agg')  # Non-interactive backend
            import matplotlib.pyplot as plt
            import seaborn as sns
            self._has_matplotlib = True
        except ImportError:
            logger.warning("Matplotlib/Seaborn not available. Visualization will be limited.")
            self._has_matplotlib = False
    
    def extract_attention_weights(
        self,
        model: Any,
        input_data: np.ndarray
    ) -> Dict[str, np.ndarray]:
        """
        Extract attention weights from a model.
        
        Args:
            model: Model with attention layers
            input_data: Input data to process
            
        Returns:
            Dict of layer names to attention weight matrices
        """
        attention_weights = {}
        
        # Try to extract from common model types
        if hasattr(model, 'attention_weights'):
            # Direct access
            attention_weights['main'] = model.attention_weights
            
        elif hasattr(model, 'get_attention_weights'):
            # Method access
            attention_weights = model.get_attention_weights(input_data)
            
        elif hasattr(model, 'layers'):
            # Iterate through layers (Keras/TF style)
            for i, layer in enumerate(model.layers):
                if 'attention' in layer.name.lower():
                    if hasattr(layer, 'attention_scores'):
                        attention_weights[f'layer_{i}_{layer.name}'] = layer.attention_scores
        
        elif hasattr(model, 'named_modules'):
            # PyTorch style
            for name, module in model.named_modules():
                if 'attention' in name.lower():
                    if hasattr(module, 'attention_weights'):
                        attention_weights[name] = module.attention_weights.detach().cpu().numpy()
        
        return attention_weights
    
    def visualize_attention_heatmap(
        self,
        attention_weights: np.ndarray,
        row_labels: List[str] = None,
        col_labels: List[str] = None,
        title: str = "Attention Weights",
        save_path: str = None
    ) -> Optional[str]:
        """
        Create a heatmap visualization of attention weights.
        
        Args:
            attention_weights: 2D attention matrix
            row_labels: Labels for rows (query tokens)
            col_labels: Labels for columns (key tokens)
            title: Plot title
            save_path: Optional path to save the figure
            
        Returns:
            Path to saved figure or None
        """
        if not self._has_matplotlib:
            logger.warning("Cannot create heatmap: matplotlib not available")
            return self._save_attention_data(attention_weights, save_path or "attention_heatmap.json")
        
        import matplotlib.pyplot as plt
        import seaborn as sns
        
        fig, ax = plt.subplots(figsize=self.config.figure_size)
        
        # Create heatmap
        sns.heatmap(
            attention_weights,
            xticklabels=col_labels if col_labels else False,
            yticklabels=row_labels if row_labels else False,
            cmap=self.config.colormap,
            annot=True if attention_weights.shape[0] <= 10 else False,
            fmt='.2f',
            ax=ax
        )
        
        ax.set_title(title)
        ax.set_xlabel('Key Tokens')
        ax.set_ylabel('Query Tokens')
        
        plt.tight_layout()
        
        # Save figure
        if save_path is None:
            save_path = self.output_dir / f"attention_heatmap_{id(attention_weights)}.png"
        else:
            save_path = Path(save_path)
        
        plt.savefig(save_path, dpi=self.config.dpi, bbox_inches='tight')
        plt.close()
        
        logger.info(f"Attention heatmap saved to {save_path}")
        return str(save_path)
    
    def visualize_multihead_attention(
        self,
        attention_weights: np.ndarray,
        head_names: List[str] = None,
        title: str = "Multi-Head Attention",
        save_path: str = None
    ) -> Optional[str]:
        """
        Visualize multi-head attention as multiple heatmaps.
        
        Args:
            attention_weights: 3D array (heads, query, key)
            head_names: Names for each attention head
            title: Plot title
            save_path: Path to save figure
            
        Returns:
            Path to saved figure
        """
        if not self._has_matplotlib:
            return self._save_attention_data(attention_weights, save_path or "multihead_attention.json")
        
        import matplotlib.pyplot as plt
        import seaborn as sns
        
        n_heads = attention_weights.shape[0]
        n_cols = min(4, n_heads)
        n_rows = (n_heads + n_cols - 1) // n_cols
        
        fig, axes = plt.subplots(n_rows, n_cols, figsize=(4 * n_cols, 4 * n_rows))
        axes = np.atleast_2d(axes)
        
        for i in range(n_heads):
            row, col = i // n_cols, i % n_cols
            ax = axes[row, col]
            
            head_name = head_names[i] if head_names and i < len(head_names) else f"Head {i+1}"
            
            sns.heatmap(
                attention_weights[i],
                cmap=self.config.colormap,
                ax=ax,
                cbar=False
            )
            ax.set_title(head_name)
        
        # Hide unused subplots
        for i in range(n_heads, n_rows * n_cols):
            row, col = i // n_cols, i % n_cols
            axes[row, col].axis('off')
        
        fig.suptitle(title, fontsize=14)
        plt.tight_layout()
        
        if save_path is None:
            save_path = self.output_dir / "multihead_attention.png"
        
        plt.savefig(save_path, dpi=self.config.dpi, bbox_inches='tight')
        plt.close()
        
        return str(save_path)
    
    def get_feature_importance_from_attention(
        self,
        attention_weights: np.ndarray,
        feature_names: List[str]
    ) -> List[Dict[str, float]]:
        """
        Extract feature importance from attention weights.
        
        Args:
            attention_weights: Attention matrix
            feature_names: Names of input features
            
        Returns:
            Sorted list of feature importance
        """
        # Average attention over all queries
        if attention_weights.ndim == 3:
            # Multi-head: average over heads first
            avg_attention = attention_weights.mean(axis=0).mean(axis=0)
        else:
            avg_attention = attention_weights.mean(axis=0)
        
        # Normalize
        avg_attention = avg_attention / avg_attention.sum() if avg_attention.sum() > 0 else avg_attention
        
        # Create importance list
        importance = []
        for i, name in enumerate(feature_names):
            if i < len(avg_attention):
                importance.append({
                    'feature': name,
                    'attention_score': float(avg_attention[i]),
                    'rank': 0
                })
        
        # Sort and assign ranks
        importance.sort(key=lambda x: x['attention_score'], reverse=True)
        for i, item in enumerate(importance):
            item['rank'] = i + 1
        
        return importance
    
    def visualize_temporal_attention(
        self,
        attention_over_time: List[np.ndarray],
        timestamps: List[str] = None,
        save_path: str = None
    ) -> Optional[str]:
        """
        Visualize how attention changes over time/sequence.
        
        Args:
            attention_over_time: List of attention matrices at each timestep
            timestamps: Labels for each timestep
            save_path: Path to save figure
            
        Returns:
            Path to saved figure
        """
        if not self._has_matplotlib:
            return None
        
        import matplotlib.pyplot as plt
        
        n_steps = len(attention_over_time)
        
        if timestamps is None:
            timestamps = [f"t={i}" for i in range(n_steps)]
        
        # Get average attention per step
        avg_attention = [w.mean() for w in attention_over_time]
        max_attention = [w.max() for w in attention_over_time]
        
        fig, ax = plt.subplots(figsize=self.config.figure_size)
        
        x = range(n_steps)
        ax.plot(x, avg_attention, 'b-o', label='Average Attention')
        ax.plot(x, max_attention, 'r-s', label='Max Attention')
        
        ax.set_xlabel('Time Step')
        ax.set_ylabel('Attention Weight')
        ax.set_title('Temporal Attention Pattern')
        ax.set_xticks(x)
        ax.set_xticklabels(timestamps, rotation=45)
        ax.legend()
        ax.grid(True, alpha=0.3)
        
        plt.tight_layout()
        
        if save_path is None:
            save_path = self.output_dir / "temporal_attention.png"
        
        plt.savefig(save_path, dpi=self.config.dpi, bbox_inches='tight')
        plt.close()
        
        return str(save_path)
    
    def visualize_attention_for_match(
        self,
        model: Any,
        match_features: Dict[str, float],
        feature_names: List[str],
        title: str = "Match Prediction Attention"
    ) -> Dict[str, Any]:
        """
        Visualize attention for a specific match prediction.
        
        Args:
            model: Model to extract attention from
            match_features: Feature dict for the match
            feature_names: Names of features
            title: Visualization title
            
        Returns:
            Dict with attention analysis
        """
        # Prepare input
        input_array = np.array([match_features.get(f, 0) for f in feature_names]).reshape(1, -1)
        
        # Extract attention
        attention = self.extract_attention_weights(model, input_array)
        
        if not attention:
            return {'error': 'Could not extract attention weights'}
        
        # Get main attention
        main_attention = list(attention.values())[0]
        
        # Get feature importance
        importance = self.get_feature_importance_from_attention(main_attention, feature_names)
        
        # Create visualization
        heatmap_path = None
        if len(main_attention.shape) == 2:
            heatmap_path = self.visualize_attention_heatmap(
                main_attention,
                row_labels=feature_names[:main_attention.shape[0]],
                col_labels=feature_names[:main_attention.shape[1]],
                title=title
            )
        
        return {
            'top_features': importance[:10],
            'attention_stats': {
                'mean': float(main_attention.mean()),
                'max': float(main_attention.max()),
                'std': float(main_attention.std())
            },
            'visualization_path': heatmap_path
        }
    
    def _save_attention_data(
        self,
        attention_weights: np.ndarray,
        filename: str
    ) -> str:
        """Save attention data as JSON when visualization not available."""
        filepath = self.output_dir / filename
        
        data = {
            'shape': list(attention_weights.shape),
            'mean': float(attention_weights.mean()),
            'max': float(attention_weights.max()),
            'min': float(attention_weights.min()),
            'data': attention_weights.tolist() if attention_weights.size < 1000 else 'Too large to save'
        }
        
        with open(filepath, 'w') as f:
            json.dump(data, f, indent=2)
        
        return str(filepath)
    
    def compare_attention_patterns(
        self,
        attention_a: np.ndarray,
        attention_b: np.ndarray,
        name_a: str = "Model A",
        name_b: str = "Model B"
    ) -> Dict[str, Any]:
        """
        Compare attention patterns between two models/predictions.
        
        Returns:
            Comparison statistics
        """
        # Correlation
        corr = np.corrcoef(attention_a.flatten(), attention_b.flatten())[0, 1]
        
        # Difference stats
        diff = attention_a - attention_b
        
        return {
            'correlation': round(float(corr), 4),
            'mean_difference': round(float(diff.mean()), 4),
            'max_difference': round(float(np.abs(diff).max()), 4),
            'similar': corr > 0.8
        }


# Global instance
_visualizer: Optional[AttentionVisualizer] = None


def get_visualizer() -> AttentionVisualizer:
    """Get or create attention visualizer."""
    global _visualizer
    if _visualizer is None:
        _visualizer = AttentionVisualizer()
    return _visualizer


def visualize_attention(
    attention_weights: np.ndarray,
    labels: List[str] = None,
    title: str = "Attention Weights"
) -> Optional[str]:
    """Quick function to visualize attention weights."""
    return get_visualizer().visualize_attention_heatmap(
        attention_weights,
        row_labels=labels,
        col_labels=labels,
        title=title
    )


def get_attention_importance(
    attention_weights: np.ndarray,
    feature_names: List[str]
) -> List[Dict]:
    """Quick function to get feature importance from attention."""
    return get_visualizer().get_feature_importance_from_attention(
        attention_weights,
        feature_names
    )