analyze_standalone.py

#!/usr/bin/env python3
"""
Standalone script for analyzing attention head explanation results.
This script only requires matplotlib, seaborn, pandas - no transformer dependencies.

Usage:
    python analyze_standalone.py
    python analyze_standalone.py --reports-dir reports
    python analyze_standalone.py --help
"""

import json
import pandas as pd
import matplotlib.pyplot as plt
import seaborn as sns
from pathlib import Path
from typing import Dict, List, Optional, Tuple
import numpy as np
import argparse
import sys


class ResultsAnalyzer:
    """Analyze and compare results from multiple experimental runs."""
    
    def __init__(self, reports_dir: str = "reports"):
        """
        Initialize the analyzer.
        
        Args:
            reports_dir: Base directory containing report subdirectories
        """
        self.reports_dir = Path(reports_dir)
        self.experiments = {}
        self._load_experiments()
    
    def _load_experiments(self):
        """Load all available experiments from the reports directory."""
        if not self.reports_dir.exists():
            print(f"Reports directory {self.reports_dir} does not exist")
            return
        
        for exp_dir in self.reports_dir.iterdir():
            if exp_dir.is_dir():
                outputs_dir = exp_dir / "outputs"
                if outputs_dir.exists():
                    self.experiments[exp_dir.name] = self._load_experiment_data(outputs_dir)
    
    def _load_experiment_data(self, outputs_dir: Path) -> Dict:
        """Load data for a single experiment."""
        data = {}
        
        # Load summary CSV
        summary_path = outputs_dir / "summary.csv"
        if summary_path.exists():
            data['summary'] = pd.read_csv(summary_path)
        
        # Load explanations
        explanations_path = outputs_dir / "explanations.jsonl"
        if explanations_path.exists():
            explanations = []
            with open(explanations_path, 'r') as f:
                for line in f:
                    explanations.append(json.loads(line))
            data['explanations'] = pd.DataFrame(explanations)
        
        # Load clusters
        clusters_path = outputs_dir / "clusters.json"
        if clusters_path.exists():
            with open(clusters_path, 'r') as f:
                data['clusters'] = json.load(f)
        
        # Load scores parquet if available
        scores_path = outputs_dir / "scores.parquet"
        if scores_path.exists():
            data['scores'] = pd.read_parquet(scores_path)
        
        return data
    
    def list_experiments(self) -> List[str]:
        """List all available experiments."""
        return list(self.experiments.keys())
    
    def get_experiment_stats(self, experiment: str) -> Dict:
        """Get summary statistics for an experiment."""
        if experiment not in self.experiments:
            raise ValueError(f"Experiment {experiment} not found")
        
        exp_data = self.experiments[experiment]
        stats = {}
        
        if 'summary' in exp_data:
            df = exp_data['summary']
            stats['n_heads'] = len(df)
            stats['n_clusters'] = df['cluster_id'].nunique() - (1 if -1 in df['cluster_id'].values else 0)
            stats['n_unclustered'] = (df['cluster_id'] == -1).sum()
            stats['mean_sim_score'] = df['sim_score'].mean()
            stats['median_sim_score'] = df['sim_score'].median()
            stats['std_sim_score'] = df['sim_score'].std()
        
        if 'clusters' in exp_data:
            clusters = exp_data['clusters']['clusters']
            stats['cluster_sizes'] = [c['n_members'] for c in clusters]
            stats['mean_agreement'] = np.mean([c['agreement'] for c in clusters])
            stats['mean_specificity'] = np.mean([c['specificity'] for c in clusters])
        
        return stats
    
    def compare_experiments(self, experiments: Optional[List[str]] = None) -> pd.DataFrame:
        """
        Compare multiple experiments side by side.
        
        Args:
            experiments: List of experiment names to compare. If None, compare all.
        
        Returns:
            DataFrame with comparison statistics
        """
        if experiments is None:
            experiments = self.list_experiments()
        
        comparison = []
        for exp in experiments:
            stats = self.get_experiment_stats(exp)
            stats['experiment'] = exp
            comparison.append(stats)
        
        return pd.DataFrame(comparison)
    
    def plot_score_distributions(self, experiments: Optional[List[str]] = None, 
                                 figsize: Tuple[int, int] = (12, 6)):
        """Plot simulation score distributions for multiple experiments."""
        if experiments is None:
            experiments = self.list_experiments()
        
        fig, axes = plt.subplots(1, 2, figsize=figsize)
        
        # Violin plot
        data_for_violin = []
        for exp in experiments:
            if 'summary' in self.experiments[exp]:
                df = self.experiments[exp]['summary']
                for score in df['sim_score']:
                    data_for_violin.append({'experiment': exp, 'sim_score': score})
        
        df_violin = pd.DataFrame(data_for_violin)
        sns.violinplot(data=df_violin, x='experiment', y='sim_score', ax=axes[0])
        axes[0].set_title('Simulation Score Distribution')
        axes[0].set_xlabel('Experiment')
        axes[0].set_ylabel('Simulation Score')
        axes[0].tick_params(axis='x', rotation=45)
        
        # Box plot
        sns.boxplot(data=df_violin, x='experiment', y='sim_score', ax=axes[1])
        axes[1].set_title('Simulation Score Box Plot')
        axes[1].set_xlabel('Experiment')
        axes[1].set_ylabel('Simulation Score')
        axes[1].tick_params(axis='x', rotation=45)
        
        plt.tight_layout()
        return fig
    
    def plot_cluster_distribution(self, experiment: str, figsize: Tuple[int, int] = (14, 8)):
        """Plot cluster distribution for a single experiment."""
        if experiment not in self.experiments:
            raise ValueError(f"Experiment {experiment} not found")
        
        exp_data = self.experiments[experiment]
        if 'summary' not in exp_data or 'clusters' not in exp_data:
            raise ValueError(f"Required data not found for {experiment}")
        
        df = exp_data['summary']
        clusters = exp_data['clusters']['clusters']
        
        fig, axes = plt.subplots(2, 2, figsize=figsize)
        
        # Cluster size distribution
        cluster_sizes = df[df['cluster_id'] != -1].groupby('cluster_id').size()
        axes[0, 0].bar(range(len(cluster_sizes)), cluster_sizes.values)
        axes[0, 0].set_title('Cluster Sizes')
        axes[0, 0].set_xlabel('Cluster ID')
        axes[0, 0].set_ylabel('Number of Heads')
        
        # Scores by cluster
        clustered_df = df[df['cluster_id'] != -1].copy()
        if len(clustered_df) > 0:
            sns.boxplot(data=clustered_df, x='cluster_id', y='sim_score', ax=axes[0, 1])
            axes[0, 1].set_title('Simulation Scores by Cluster')
            axes[0, 1].set_xlabel('Cluster ID')
            axes[0, 1].set_ylabel('Simulation Score')
            axes[0, 1].tick_params(axis='x', rotation=45)
        
        # Cluster agreement scores
        cluster_ids = [c['cluster_id'] for c in clusters]
        agreements = [c['agreement'] for c in clusters]
        axes[1, 0].bar(range(len(cluster_ids)), agreements)
        axes[1, 0].set_title('Cluster Agreement Scores')
        axes[1, 0].set_xlabel('Cluster Index')
        axes[1, 0].set_ylabel('Agreement')
        axes[1, 0].axhline(y=0.8, color='r', linestyle='--', label='0.8 threshold')
        axes[1, 0].legend()
        
        # Top clusters by size
        top_clusters = sorted(clusters, key=lambda x: x['n_members'], reverse=True)[:10]
        labels = [c['canonical_label'][:30] + '...' if len(c['canonical_label']) > 30 
                 else c['canonical_label'] for c in top_clusters]
        sizes = [c['n_members'] for c in top_clusters]
        
        axes[1, 1].barh(range(len(labels)), sizes)
        axes[1, 1].set_yticks(range(len(labels)))
        axes[1, 1].set_yticklabels(labels, fontsize=8)
        axes[1, 1].set_title('Top 10 Clusters by Size')
        axes[1, 1].set_xlabel('Number of Heads')
        axes[1, 1].invert_yaxis()
        
        plt.suptitle(f'Cluster Analysis: {experiment}', fontsize=14, y=1.00)
        plt.tight_layout()
        return fig
    
    def plot_layer_head_heatmap(self, experiment: str, metric: str = 'sim_score',
                               figsize: Tuple[int, int] = (12, 8)):
        """Plot a heatmap of scores across layers and heads."""
        if experiment not in self.experiments:
            raise ValueError(f"Experiment {experiment} not found")
        
        df = self.experiments[experiment]['summary']
        
        # Create pivot table
        pivot = df.pivot(index='layer', columns='head', values=metric)
        
        fig, ax = plt.subplots(figsize=figsize)
        
        if metric == 'sim_score':
            sns.heatmap(pivot, annot=True, fmt='.2f', cmap='YlOrRd', ax=ax, cbar_kws={'label': 'Simulation Score'})
            title = f'Simulation Scores Heatmap: {experiment}'
        else:
            sns.heatmap(pivot, annot=True, fmt='.0f', cmap='tab20', ax=ax, cbar_kws={'label': 'Cluster ID'})
            title = f'Cluster Assignment Heatmap: {experiment}'
        
        ax.set_title(title)
        ax.set_xlabel('Head')
        ax.set_ylabel('Layer')
        
        plt.tight_layout()
        return fig
    
    def plot_top_heads(self, experiment: str, n: int = 20, figsize: Tuple[int, int] = (12, 8)):
        """Plot the top N heads by simulation score."""
        if experiment not in self.experiments:
            raise ValueError(f"Experiment {experiment} not found")
        
        df = self.experiments[experiment]['summary']
        top_heads = df.nlargest(n, 'sim_score')
        
        fig, ax = plt.subplots(figsize=figsize)
        
        # Create labels
        labels = [f"L{row['layer']}H{row['head']}" for _, row in top_heads.iterrows()]
        scores = top_heads['sim_score'].values
        colors = ['red' if row['cluster_id'] == -1 else 'blue' 
                 for _, row in top_heads.iterrows()]
        
        bars = ax.barh(range(len(labels)), scores, color=colors, alpha=0.7)
        ax.set_yticks(range(len(labels)))
        ax.set_yticklabels(labels)
        ax.set_xlabel('Simulation Score')
        ax.set_title(f'Top {n} Attention Heads by Simulation Score: {experiment}')
        ax.invert_yaxis()
        
        # Add legend
        from matplotlib.patches import Patch
        legend_elements = [Patch(facecolor='blue', alpha=0.7, label='Clustered'),
                          Patch(facecolor='red', alpha=0.7, label='Unclustered')]
        ax.legend(handles=legend_elements)
        
        plt.tight_layout()
        return fig
    
    def plot_comparison_summary(self, experiments: Optional[List[str]] = None,
                               figsize: Tuple[int, int] = (16, 10)):
        """Create a comprehensive comparison plot across experiments."""
        if experiments is None:
            experiments = self.list_experiments()
        
        fig, axes = plt.subplots(2, 3, figsize=figsize)
        
        stats_data = []
        for exp in experiments:
            stats = self.get_experiment_stats(exp)
            stats['experiment'] = exp
            stats_data.append(stats)
        
        df_stats = pd.DataFrame(stats_data)
        
        # Number of heads
        axes[0, 0].bar(df_stats['experiment'], df_stats['n_heads'])
        axes[0, 0].set_title('Number of Heads Analyzed')
        axes[0, 0].set_ylabel('Count')
        axes[0, 0].tick_params(axis='x', rotation=45)
        
        # Number of clusters
        axes[0, 1].bar(df_stats['experiment'], df_stats['n_clusters'])
        axes[0, 1].set_title('Number of Clusters Found')
        axes[0, 1].set_ylabel('Count')
        axes[0, 1].tick_params(axis='x', rotation=45)
        
        # Unclustered heads
        axes[0, 2].bar(df_stats['experiment'], df_stats['n_unclustered'], color='orange')
        axes[0, 2].set_title('Unclustered Heads')
        axes[0, 2].set_ylabel('Count')
        axes[0, 2].tick_params(axis='x', rotation=45)
        
        # Mean simulation scores
        axes[1, 0].bar(df_stats['experiment'], df_stats['mean_sim_score'], color='green')
        axes[1, 0].set_title('Mean Simulation Score')
        axes[1, 0].set_ylabel('Score')
        axes[1, 0].tick_params(axis='x', rotation=45)
        
        # Mean agreement
        if 'mean_agreement' in df_stats.columns:
            axes[1, 1].bar(df_stats['experiment'], df_stats['mean_agreement'], color='purple')
            axes[1, 1].set_title('Mean Cluster Agreement')
            axes[1, 1].set_ylabel('Agreement')
            axes[1, 1].axhline(y=0.8, color='r', linestyle='--', alpha=0.5)
            axes[1, 1].tick_params(axis='x', rotation=45)
        
        # Mean specificity
        if 'mean_specificity' in df_stats.columns:
            axes[1, 2].bar(df_stats['experiment'], df_stats['mean_specificity'], color='teal')
            axes[1, 2].set_title('Mean Cluster Specificity')
            axes[1, 2].set_ylabel('Specificity')
            axes[1, 2].tick_params(axis='x', rotation=45)
        
        plt.suptitle('Experiment Comparison Summary', fontsize=16, y=1.00)
        plt.tight_layout()
        return fig


def main():
    parser = argparse.ArgumentParser(
        description="Analyze attention head explanation results",
        formatter_class=argparse.ArgumentDefaultsHelpFormatter
    )
    parser.add_argument(
        "--reports-dir",
        type=str,
        default="reports",
        help="Directory containing experiment reports"
    )
    parser.add_argument(
        "--output-dir",
        type=str,
        default=None,
        help="Output directory for plots (default: same as reports-dir)"
    )
    
    args = parser.parse_args()
    
    # Use non-interactive backend for server environments
    import matplotlib
    matplotlib.use('Agg')
    
    print("=" * 80)
    print("Attention Head Analysis - Results Visualization")
    print("=" * 80)
    
    # Initialize analyzer
    print(f"\nLoading experiments from: {args.reports_dir}")
    analyzer = ResultsAnalyzer(reports_dir=args.reports_dir)
    
    # List experiments
    experiments = analyzer.list_experiments()
    if not experiments:
        print(f"\nNo experiments found in {args.reports_dir}")
        print("\nExpected structure:")
        print(f"  {args.reports_dir}/")
        print("    experiment-name/")
        print("      outputs/")
        print("        summary.csv")
        print("        explanations.jsonl")
        print("        clusters.json")
        print("        ...")
        sys.exit(1)
    
    print(f"\nFound {len(experiments)} experiment(s):")
    for exp in experiments:
        print(f"  - {exp}")
    
    # Get comparison statistics
    print("\n" + "=" * 80)
    print("Experiment Comparison")
    print("=" * 80)
    comparison = analyzer.compare_experiments()
    print(comparison.to_string(index=False))
    
    # Set output directory
    output_dir = Path(args.output_dir) if args.output_dir else Path(args.reports_dir)
    
    # Generate plots for each experiment
    print("\n" + "=" * 80)
    print("Generating Visualizations")
    print("=" * 80)
    
    for exp in experiments:
        print(f"\nProcessing {exp}...")
        exp_output_dir = Path(args.reports_dir) / exp
        exp_output_dir.mkdir(exist_ok=True, parents=True)
        
        try:
            # Cluster distribution
            print("  - Cluster analysis...")
            fig = analyzer.plot_cluster_distribution(exp)
            fig.savefig(exp_output_dir / 'cluster_analysis.png', dpi=150, bbox_inches='tight')
            plt.close(fig)
            
            # Heatmaps
            print("  - Simulation score heatmap...")
            fig = analyzer.plot_layer_head_heatmap(exp, 'sim_score')
            fig.savefig(exp_output_dir / 'heatmap_scores.png', dpi=150, bbox_inches='tight')
            plt.close(fig)
            
            print("  - Cluster assignment heatmap...")
            fig = analyzer.plot_layer_head_heatmap(exp, 'cluster_id')
            fig.savefig(exp_output_dir / 'heatmap_clusters.png', dpi=150, bbox_inches='tight')
            plt.close(fig)
            
            # Top heads
            print("  - Top heads...")
            fig = analyzer.plot_top_heads(exp, n=20)
            fig.savefig(exp_output_dir / 'top_heads.png', dpi=150, bbox_inches='tight')
            plt.close(fig)
            
            print(f"  ✓ Plots saved to {exp_output_dir}/")
            
        except Exception as e:
            print(f"  ✗ Error generating plots: {e}")
            import traceback
            traceback.print_exc()
    
    # Generate comparison plots if multiple experiments
    if len(experiments) > 1:
        print(f"\nGenerating comparison plots...")
        
        try:
            print("  - Score distributions...")
            fig = analyzer.plot_score_distributions(experiments)
            fig.savefig(output_dir / 'score_comparison.png', dpi=150, bbox_inches='tight')
            plt.close(fig)
            
            print("  - Experiment comparison summary...")
            fig = analyzer.plot_comparison_summary(experiments)
            fig.savefig(output_dir / 'experiment_comparison.png', dpi=150, bbox_inches='tight')
            plt.close(fig)
            
            print(f"  ✓ Comparison plots saved to {output_dir}")
            
        except Exception as e:
            print(f"  ✗ Error generating comparison plots: {e}")
            import traceback
            traceback.print_exc()
    
    print("\n" + "=" * 80)
    print("Analysis Complete!")
    print("=" * 80)
    print(f"\nResults saved to: {output_dir}")
    print("\nGenerated plots:")
    for exp in experiments:
        exp_dir = Path(args.reports_dir) / exp
        print(f"  {exp}:")
        for plot_file in ['cluster_analysis.png', 'heatmap_scores.png', 
                         'heatmap_clusters.png', 'top_heads.png']:
            plot_path = exp_dir / plot_file
            if plot_path.exists():
                print(f"    - {plot_path}")
    
    if len(experiments) > 1:
        print(f"  Comparisons:")
        for plot_file in ['score_comparison.png', 'experiment_comparison.png']:
            plot_path = output_dir / plot_file
            if plot_path.exists():
                print(f"    - {plot_path}")
    
    print("=" * 80)


if __name__ == "__main__":
    main()