xai/visualization.py

"""Visualization utilities for XAI results."""

import json
from pathlib import Path
from typing import Any, Dict, List, Optional, Union

import matplotlib.pyplot as plt
import numpy as np

try:
    import shap
    HAS_SHAP = True
except ImportError:
    HAS_SHAP = False

try:
    from IPython.display import HTML, display
    HAS_IPYTHON = True
except ImportError:
    HAS_IPYTHON = False


class XAIVisualizer:
    """Visualize XAI results."""
    
    def __init__(
        self,
        class_names: Optional[List[str]] = None,
        output_dir: str = "outputs/xai",
    ):
        self.class_names = class_names or [
            "negative", "neutral", "positive", "sarcastic"
        ]
        self.output_dir = Path(output_dir)
        self.output_dir.mkdir(parents=True, exist_ok=True)
    
    def plot_word_importance(
        self,
        words: List[str],
        importance: List[float],
        title: str = "Word Importance",
        output_path: Optional[str] = None,
        figsize: tuple = (10, 6),
    ) -> plt.Figure:
        """Plot word importance as horizontal bar chart.
        
        Args:
            words: List of words/tokens
            importance: Importance scores
            title: Plot title
            output_path: Path to save figure
            figsize: Figure size
        
        Returns:
            Matplotlib figure
        """
        fig, ax = plt.subplots(figsize=figsize)
        
        # Sort by absolute importance
        sorted_pairs = sorted(
            zip(words, importance),
            key=lambda x: abs(x[1]),
            reverse=True,
        )
        
        sorted_words = [p[0] for p in sorted_pairs]
        sorted_importance = [p[1] for p in sorted_pairs]
        
        # Color based on positive/negative
        colors = ["green" if v > 0 else "red" for v in sorted_importance]
        
        ax.barh(sorted_words, sorted_importance, color=colors, alpha=0.7)
        ax.axvline(x=0, color="black", linestyle="-", linewidth=0.5)
        ax.set_xlabel("SHAP Value")
        ax.set_title(title)
        ax.invert_yaxis()
        
        plt.tight_layout()
        
        if output_path:
            plt.savefig(output_path, dpi=150, bbox_inches="tight")
        
        return fig
    
    def plot_feature_importance(
        self,
        features: List[str],
        importance: List[float],
        title: str = "Feature Importance",
        output_path: Optional[str] = None,
        figsize: tuple = (10, 6),
    ) -> plt.Figure:
        """Plot feature importance bar chart.
        
        Args:
            features: List of feature names
            importance: Importance scores
            title: Plot title
            output_path: Path to save figure
            figsize: Figure size
        
        Returns:
            Matplotlib figure
        """
        fig, ax = plt.subplots(figsize=figsize)
        
        # Sort by importance
        sorted_pairs = sorted(
            zip(features, importance),
            key=lambda x: x[1],
            reverse=True,
        )
        
        sorted_features = [p[0] for p in sorted_pairs]
        sorted_importance = [p[1] for p in sorted_pairs]
        
        ax.barh(sorted_features, sorted_importance, color="steelblue", alpha=0.7)
        ax.set_xlabel("Importance Score")
        ax.set_title(title)
        ax.invert_yaxis()
        
        plt.tight_layout()
        
        if output_path:
            plt.savefig(output_path, dpi=150, bbox_inches="tight")
        
        return fig
    
    def plot_confidence_distribution(
        self,
        predictions: List[str],
        confidences: List[float],
        output_path: Optional[str] = None,
        figsize: tuple = (10, 6),
    ) -> plt.Figure:
        """Plot distribution of prediction confidences.
        
        Args:
            predictions: Predicted classes
            confidences: Confidence scores
            output_path: Path to save figure
            figsize: Figure size
        
        Returns:
            Matplotlib figure
        """
        fig, (ax1, ax2) = plt.subplots(1, 2, figsize=figsize)
        
        # Histogram
        ax1.hist(confidences, bins=20, alpha=0.7, color="steelblue")
        ax1.axvline(np.mean(confidences), color="red", linestyle="--", 
                    label=f"Mean: {np.mean(confidences):.3f}")
        ax1.set_xlabel("Confidence")
        ax1.set_ylabel("Count")
        ax1.set_title("Confidence Distribution")
        ax1.legend()
        
        # By class
        unique_classes = list(set(predictions))
        class_confidences = {c: [] for c in unique_classes}
        
        for pred, conf in zip(predictions, confidences):
            class_confidences[pred].append(conf)
        
        ax2.boxplot(
            [class_confidences[c] for c in unique_classes],
            labels=unique_classes,
        )
        ax2.set_xlabel("Class")
        ax2.set_ylabel("Confidence")
        ax2.set_title("Confidence by Class")
        
        plt.tight_layout()
        
        if output_path:
            plt.savefig(output_path, dpi=150, bbox_inches="tight")
        
        return fig
    
    def plot_shap_summary(
        self,
        shap_values: np.ndarray,
        features: np.ndarray,
        feature_names: List[str],
        output_path: Optional[str] = None,
        figsize: tuple = (12, 8),
    ) -> plt.Figure:
        """Plot SHAP summary (beeswarm) plot.
        
        Args:
            shap_values: SHAP values array
            features: Feature values array
            feature_names: Names of features
            output_path: Path to save figure
            figsize: Figure size
        
        Returns:
            Matplotlib figure
        """
        if not HAS_SHAP:
            raise ImportError("shap library required for this visualization")
        
        fig, ax = plt.subplots(figsize=figsize)
        
        shap.summary_plot(
            shap_values,
            features,
            feature_names=feature_names,
            show=False,
        )
        
        plt.tight_layout()
        
        if output_path:
            plt.savefig(output_path, dpi=150, bbox_inches="tight")
        
        return fig
    
    def plot_comparison(
        self,
        explanations: Dict[str, List[Tuple[str, float]]],
        output_path: Optional[str] = None,
        figsize: tuple = (12, 8),
    ) -> plt.Figure:
        """Compare explanations across different methods or samples.
        
        Args:
            explanations: Dict mapping sample IDs to explanation tuples
            output_path: Path to save figure
            figsize: Figure size
        
        Returns:
            Matplotlib figure
        """
        fig, ax = plt.subplots(figsize=figsize)
        
        # Get all unique words
        all_words = set()
        for exp in explanations.values():
            for word, _ in exp:
                all_words.add(word)
        
        all_words = list(all_words)[:20]  # Limit to top 20
        
        # Create matrix
        matrix = []
        for sample_id, exp in explanations.items():
            exp_dict = dict(exp)
            row = [exp_dict.get(w, 0) for w in all_words]
            matrix.append(row)
        
        matrix = np.array(matrix)
        
        # Plot heatmap
        im = ax.imshow(matrix, cmap="RdBu_r", aspect="auto")
        ax.set_xticks(range(len(all_words)))
        ax.set_xticklabels(all_words, rotation=45, ha="right")
        ax.set_yticks(range(len(explanations)))
        ax.set_yticklabels(list(explanations.keys()))
        ax.set_title("Explanation Comparison")
        
        plt.colorbar(im, ax=ax, label="Importance")
        
        plt.tight_layout()
        
        if output_path:
            plt.savefig(output_path, dpi=150, bbox_inches="tight")
        
        return fig
    
    def save_explanations(
        self,
        explanations: List[Dict[str, Any]],
        output_path: str,
    ) -> None:
        """Save explanations to JSON file.
        
        Args:
            explanations: List of explanation dictionaries
            output_path: Path to save JSON
        """
        with open(output_path, "w", encoding="utf-8") as f:
            json.dump(explanations, f, indent=2, ensure_ascii=False)
    
    def generate_html_report(
        self,
        explanations: List[Dict[str, Any]],
        output_path: str,
    ) -> None:
        """Generate HTML report of explanations.
        
        Args:
            explanations: List of explanation dictionaries
            output_path: Path to save HTML
        """
        html = """
        <!DOCTYPE html>
        <html>
        <head>
            <title>XAI Explanation Report</title>
            <style>
                body { font-family: Arial, sans-serif; margin: 20px; }
                .explanation { border: 1px solid #ccc; padding: 15px; margin: 10px 0; border-radius: 5px; }
                .text { font-size: 18px; margin-bottom: 10px; }
                .prediction { font-weight: bold; color: #2196F3; }
                .word { display: inline-block; padding: 2px 5px; margin: 2px; border-radius: 3px; }
                .positive { background-color: #c8e6c9; }
                .negative { background-color: #ffcdd2; }
                .neutral { background-color: #e0e0e0; }
                table { border-collapse: collapse; width: 100%; }
                th, td { border: 1px solid #ddd; padding: 8px; text-align: left; }
                th { background-color: #f5f5f5; }
            </style>
        </head>
        <body>
            <h1>XAI Explanation Report</h1>
            <p>Total explanations: """ + str(len(explanations)) + """</p>
        """
        
        for i, exp in enumerate(explanations):
            html += f"""
            <div class="explanation">
                <div class="text">{exp.get('text', 'N/A')}</div>
                <div class="prediction">Predicted: {exp.get('predicted_class', 'N/A')}</div>
                <div>
            """
            
            for word, weight in exp.get("word_importance", []):
                color_class = "positive" if weight > 0 else "negative"
                html += f'<span class="word {color_class}">{word}: {weight:.3f}</span>'
            
            html += """
                </div>
            </div>
            """
        
        html += """
        </body>
        </html>
        """
        
        with open(output_path, "w", encoding="utf-8") as f:
            f.write(html)


def create_visualizer(
    class_names: Optional[List[str]] = None,
    output_dir: str = "outputs/xai",
) -> XAIVisualizer:
    """Factory function to create XAI visualizer."""
    return XAIVisualizer(
        class_names=class_names,
        output_dir=output_dir,
    )


if __name__ == "__main__":
    visualizer = create_visualizer()
    print("XAIVisualizer loaded")
    print(f"Available methods: plot_word_importance, plot_feature_importance, etc.")