Blame for code/pmps/analysis/plot_collapse_analysis.py

"""
Collapse Analysis Visualization (PMPS vs CODI)

Generates publication-quality figures for latent representation collapse analysis:
  Fig 1: Pairwise cosine similarity heatmaps (side-by-side CODI vs PMPS)
  Fig 2: Off-diagonal cosine similarity distribution (violin plot)
  Fig 3: Singular value spectrum comparison
  Fig 4: Effective rank distribution comparison
  Fig 5: t-SNE / PCA visualization of latent embeddings

Usage:
    python analysis/plot_collapse_analysis.py \
        --pmps_dir analysis/results/pmps \
        --codi_dir analysis/results/codi \
        --output_dir analysis/figures/collapse
"""

import os
import json
import argparse
import numpy as np
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
import matplotlib.gridspec as gridspec
from matplotlib.colors import LinearSegmentedColormap
from sklearn.decomposition import PCA
from sklearn.manifold import TSNE

# ─────────────────────────────────────────────────────────────────────────────
# Global style
# ─────────────────────────────────────────────────────────────────────────────

plt.rcParams.update({
    'font.family': 'serif',
    'font.serif': ['Times New Roman', 'DejaVu Serif'],
    'font.size': 11,
    'axes.titlesize': 13,
    'axes.labelsize': 12,
    'xtick.labelsize': 10,
    'ytick.labelsize': 10,
    'legend.fontsize': 10,
    'figure.dpi': 300,
    'savefig.bbox': 'tight',
    'savefig.pad_inches': 0.05,
    'axes.linewidth': 0.8,
    'axes.grid': False,
})

COLORS = {
    'codi': '#E74C3C',
    'pmps': '#2E86C1',
    'simcot': '#27AE60',
    'cot': '#8E44AD',
}

METHOD_LABELS = {
    'codi': 'CODI',
    'pmps': 'PMPS (Ours)',
    'simcot': 'Sim-CoT',
}


def load_results(result_dir):
    """Load extracted metrics from a result directory."""
    metrics_path = os.path.join(result_dir, 'metrics.json')
    latent_path = os.path.join(result_dir, 'latent_hiddens.npy')

    with open(metrics_path, 'r') as f:
        metrics = json.load(f)

    latent_hiddens = np.load(latent_path)  # (num_samples, N_l, D)

    return metrics, latent_hiddens


# ─────────────────────────────────────────────────────────────────────────────
# Fig 1: Pairwise Cosine Similarity Heatmaps
# ─────────────────────────────────────────────────────────────────────────────

def plot_similarity_heatmaps(metrics_dict, output_path, methods=None):
    """Side-by-side average pairwise cosine similarity heatmaps."""
    if methods is None:
        methods = list(metrics_dict.keys())

    n_methods = len(methods)
    fig, axes = plt.subplots(1, n_methods, figsize=(4.5 * n_methods, 4))
    if n_methods == 1:
        axes = [axes]

    cmap = LinearSegmentedColormap.from_list(
        'sim', ['#FFFFFF', '#AED6F1', '#3498DB', '#1A5276'], N=256
    )

    for idx, method in enumerate(methods):
        metrics = metrics_dict[method]
        sim_matrices = [np.array(m['cos_sim_matrix']) for m in metrics]
        N_l = sim_matrices[0].shape[0]
        avg_matrix = np.mean(sim_matrices, axis=0)

        ax = axes[idx]
        im = ax.imshow(avg_matrix, cmap=cmap, vmin=0.0, vmax=1.0, aspect='equal')

        for i in range(N_l):
            for j in range(N_l):
                color = 'white' if avg_matrix[i, j] > 0.7 else 'black'
                ax.text(j, i, f'{avg_matrix[i, j]:.2f}',
                        ha='center', va='center', fontsize=9, color=color)

        ax.set_title(METHOD_LABELS.get(method, method), fontsize=14, fontweight='bold')
        ax.set_xlabel('Latent Token Index')
        ax.set_ylabel('Latent Token Index')
        ax.set_xticks(range(N_l))
        ax.set_yticks(range(N_l))
        ax.set_xticklabels([f'$l_{{{i+1}}}$' for i in range(N_l)])
        ax.set_yticklabels([f'$l_{{{i+1}}}$' for i in range(N_l)])

    fig.colorbar(im, ax=axes, label='Cosine Similarity', shrink=0.8, pad=0.02)
    plt.savefig(output_path, dpi=300)
    plt.savefig(output_path.replace('.png', '.pdf'), dpi=300)
    plt.close()
    print(f'  Saved: {output_path}')


# ─────────────────────────────────────────────────────────────────────────────
# Fig 2: Off-diagonal Cosine Similarity Distribution
# ─────────────────────────────────────────────────────────────────────────────

def plot_similarity_distribution(metrics_dict, output_path, methods=None):
    """Violin + box plot of off-diagonal cosine similarities."""
    if methods is None:
        methods = list(metrics_dict.keys())

    fig, ax = plt.subplots(figsize=(3 + 1.5 * len(methods), 4.5))

    data_list = []
    labels = []
    colors = []

    for method in methods:
        metrics = metrics_dict[method]
        off_diag_sims = []
        for m in metrics:
            sim_matrix = np.array(m['cos_sim_matrix'])
            N_l = sim_matrix.shape[0]
            mask = ~np.eye(N_l, dtype=bool)
            off_diag_sims.extend(sim_matrix[mask].tolist())
        data_list.append(off_diag_sims)
        labels.append(METHOD_LABELS.get(method, method))
        colors.append(COLORS.get(method, '#888888'))

    positions = range(1, len(methods) + 1)
    parts = ax.violinplot(data_list, positions=positions, showmedians=False, showextrema=False)
    for pc, color in zip(parts['bodies'], colors):
        pc.set_facecolor(color)
        pc.set_alpha(0.3)

    bp = ax.boxplot(data_list, positions=positions, widths=0.15,
                    patch_artist=True, showfliers=False,
                    medianprops={'color': 'black', 'linewidth': 2})
    for patch, color in zip(bp['boxes'], colors):
        patch.set_facecolor(color)
        patch.set_alpha(0.8)

    for i, (d, color) in enumerate(zip(data_list, colors)):
        mean_val = np.mean(d)
        ax.plot(i + 1, mean_val, 'D', color='white', markeredgecolor=color,
                markersize=7, markeredgewidth=1.5, zorder=5)
        ax.annotate(f'{mean_val:.3f}', (i + 1, mean_val),
                    textcoords='offset points', xytext=(15, 5),
                    fontsize=9, color=color, fontweight='bold')

    ax.set_xticks(positions)
    ax.set_xticklabels(labels, fontsize=12)
    ax.set_ylabel('Off-Diagonal Cosine Similarity', fontsize=12)
    ax.set_title('Latent Token Pairwise Similarity\n(Lower = Less Collapse)', fontsize=13)
    ax.axhline(y=1.0, color='gray', linestyle='--', alpha=0.4, label='Perfect Collapse')
    ax.legend(loc='lower right', fontsize=9)

    plt.tight_layout()
    plt.savefig(output_path, dpi=300)
    plt.savefig(output_path.replace('.png', '.pdf'), dpi=300)
    plt.close()
    print(f'  Saved: {output_path}')


# ─────────────────────────────────────────────────────────────────────────────
# Fig 3: Singular Value Spectrum
# ─────────────────────────────────────────────────────────────────────────────

def plot_singular_values(metrics_dict, output_path, methods=None):
    """Normalized singular value spectrum comparison."""
    if methods is None:
        methods = list(metrics_dict.keys())

    fig, axes = plt.subplots(1, 2, figsize=(11, 4.5))

    # Left: Normalized singular value spectrum (bar chart)
    ax = axes[0]
    bar_width = 0.8 / len(methods)
    for m_idx, method in enumerate(methods):
        metrics = metrics_dict[method]
        all_sv_norm = [m['singular_values_normalized'] for m in metrics]
        N_l = len(all_sv_norm[0])
        avg_sv = np.mean(all_sv_norm, axis=0)
        std_sv = np.std(all_sv_norm, axis=0)

        x = np.arange(N_l) + m_idx * bar_width
        ax.bar(x, avg_sv, width=bar_width, yerr=std_sv, capsize=3,
               color=COLORS.get(method, '#888'), alpha=0.8,
               label=METHOD_LABELS.get(method, method))

    ax.set_xlabel('Singular Value Index')
    ax.set_ylabel('Normalized Magnitude ($\\sigma_i / \\sum \\sigma_j$)')
    ax.set_title('(a) Singular Value Spectrum')
    ax.set_xticks(np.arange(N_l) + bar_width * (len(methods) - 1) / 2)
    ax.set_xticklabels([f'$\\sigma_{{{i+1}}}$' for i in range(N_l)])
    ax.legend()

    # Right: Cumulative explained variance
    ax2 = axes[1]
    for method in methods:
        metrics = metrics_dict[method]
        all_cum_var = [m['cumulative_variance'] for m in metrics]
        avg_cum = np.mean(all_cum_var, axis=0)
        std_cum = np.std(all_cum_var, axis=0)
        x = np.arange(1, len(avg_cum) + 1)
        ax2.plot(x, avg_cum, 'o-', color=COLORS.get(method, '#888'),
                 linewidth=2, markersize=7, label=METHOD_LABELS.get(method, method))
        ax2.fill_between(x, avg_cum - std_cum, np.minimum(avg_cum + std_cum, 1.0),
                         alpha=0.15, color=COLORS.get(method, '#888'))

    ax2.axhline(y=0.95, color='gray', linestyle='--', alpha=0.5)
    ax2.text(0.5, 0.96, '95%', color='gray', fontsize=9)
    ax2.set_xlabel('Number of Components')
    ax2.set_ylabel('Cumulative Explained Variance')
    ax2.set_title('(b) Cumulative Variance')
    ax2.set_xticks(range(1, N_l + 1))
    ax2.legend()

    plt.tight_layout()
    plt.savefig(output_path, dpi=300)
    plt.savefig(output_path.replace('.png', '.pdf'), dpi=300)
    plt.close()
    print(f'  Saved: {output_path}')


# ─────────────────────────────────────────────────────────────────────────────
# Fig 4: Effective Rank Distribution
# ─────────────────────────────────────────────────────────────────────────────

def plot_effective_rank(metrics_dict, output_path, methods=None):
    """Effective rank comparison bar chart with individual sample scatter."""
    if methods is None:
        methods = list(metrics_dict.keys())

    fig, ax = plt.subplots(figsize=(3 + 1.2 * len(methods), 4.5))

    for m_idx, method in enumerate(methods):
        metrics = metrics_dict[method]
        ranks = [m['effective_rank'] for m in metrics]
        mean_rank = np.mean(ranks)
        std_rank = np.std(ranks)
        color = COLORS.get(method, '#888')

        ax.bar(m_idx, mean_rank, width=0.5, yerr=std_rank, capsize=5,
               color=color, alpha=0.7, label=METHOD_LABELS.get(method, method))

        jitter = np.random.normal(0, 0.05, len(ranks))
        ax.scatter(m_idx + jitter, ranks, color=color, alpha=0.15, s=10, zorder=3)

        ax.text(m_idx, mean_rank + std_rank + 0.1, f'{mean_rank:.2f}',
                ha='center', fontsize=11, fontweight='bold', color=color)

    N_l = len(metrics_dict[methods[0]][0]['cos_sim_matrix'])
    ax.axhline(y=N_l, color='gray', linestyle='--', alpha=0.4)
    ax.text(len(methods) - 0.5, N_l + 0.05, f'Max = {N_l}', color='gray', fontsize=9)
    ax.axhline(y=1.0, color='red', linestyle='--', alpha=0.3)
    ax.text(len(methods) - 0.5, 1.05, 'Full Collapse', color='red', fontsize=9)

    ax.set_xticks(range(len(methods)))
    ax.set_xticklabels([METHOD_LABELS.get(m, m) for m in methods], fontsize=12)
    ax.set_ylabel('Effective Rank', fontsize=12)
    ax.set_title('Effective Rank of Latent Embeddings\n(Higher = More Diverse)', fontsize=13)
    ax.set_ylim(0, N_l + 1)

    plt.tight_layout()
    plt.savefig(output_path, dpi=300)
    plt.savefig(output_path.replace('.png', '.pdf'), dpi=300)
    plt.close()
    print(f'  Saved: {output_path}')


# ─────────────────────────────────────────────────────────────────────────────
# Fig 5: t-SNE / PCA Visualization
# ─────────────────────────────────────────────────────────────────────────────

def plot_tsne_visualization(hiddens_dict, output_path, methods=None,
                            max_samples=100, perplexity=30):
    """t-SNE visualization of latent embeddings colored by latent position."""
    if methods is None:
        methods = list(hiddens_dict.keys())

    n_methods = len(methods)
    fig, axes = plt.subplots(1, n_methods, figsize=(5.5 * n_methods, 5))
    if n_methods == 1:
        axes = [axes]

    all_embeddings = []
    all_labels = []
    all_method_labels = []

    for method in methods:
        hiddens = hiddens_dict[method][:max_samples]  # (S, N_l, D)
        S, N_l, D = hiddens.shape
        flat = hiddens.reshape(-1, D)
        positions = np.tile(np.arange(N_l), S)
        method_ids = np.full(len(positions), methods.index(method))
        all_embeddings.append(flat)
        all_labels.append(positions)
        all_method_labels.append(method_ids)

    combined = np.concatenate(all_embeddings, axis=0)
    combined_labels = np.concatenate(all_labels, axis=0)
    combined_methods = np.concatenate(all_method_labels, axis=0)

    # Fit PCA first to reduce dimensionality, then t-SNE
    pca = PCA(n_components=min(50, combined.shape[1]))
    combined_pca = pca.fit_transform(combined)

    tsne = TSNE(n_components=2, perplexity=perplexity, random_state=42, max_iter=1000)
    combined_2d = tsne.fit_transform(combined_pca)

    N_l = hiddens_dict[methods[0]].shape[1]
    cmap = plt.cm.get_cmap('tab10', N_l)

    offset = 0
    for m_idx, method in enumerate(methods):
        n_points = hiddens_dict[method][:max_samples].shape[0] * N_l
        points = combined_2d[offset:offset + n_points]
        labels = combined_labels[offset:offset + n_points]
        offset += n_points

        ax = axes[m_idx]
        for lat_idx in range(N_l):
            mask = labels == lat_idx
            ax.scatter(points[mask, 0], points[mask, 1],
                       c=[cmap(lat_idx)], s=15, alpha=0.5,
                       label=f'$l_{{{lat_idx+1}}}$')

        ax.set_title(METHOD_LABELS.get(method, method), fontsize=14, fontweight='bold')
        ax.set_xlabel('t-SNE 1')
        ax.set_ylabel('t-SNE 2')
        if m_idx == n_methods - 1:
            ax.legend(bbox_to_anchor=(1.02, 1), loc='upper left',
                      fontsize=9, markerscale=1.5)

    plt.tight_layout()
    plt.savefig(output_path, dpi=300)
    plt.savefig(output_path.replace('.png', '.pdf'), dpi=300)
    plt.close()
    print(f'  Saved: {output_path}')


def plot_pca_visualization(hiddens_dict, output_path, methods=None, max_samples=100):
    """PCA visualization — same idea as t-SNE but deterministic and faster."""
    if methods is None:
        methods = list(hiddens_dict.keys())

    n_methods = len(methods)
    fig, axes = plt.subplots(1, n_methods, figsize=(5.5 * n_methods, 5))
    if n_methods == 1:
        axes = [axes]

    all_embeddings = []
    for method in methods:
        hiddens = hiddens_dict[method][:max_samples]
        flat = hiddens.reshape(-1, hiddens.shape[-1])
        all_embeddings.append(flat)

    combined = np.concatenate(all_embeddings, axis=0)
    pca = PCA(n_components=2)
    combined_2d = pca.fit_transform(combined)

    N_l = hiddens_dict[methods[0]].shape[1]
    cmap = plt.cm.get_cmap('tab10', N_l)

    offset = 0
    for m_idx, method in enumerate(methods):
        S = min(max_samples, hiddens_dict[method].shape[0])
        n_points = S * N_l
        points = combined_2d[offset:offset + n_points]
        offset += n_points

        labels = np.tile(np.arange(N_l), S)

        ax = axes[m_idx]
        for lat_idx in range(N_l):
            mask = labels == lat_idx
            ax.scatter(points[mask, 0], points[mask, 1],
                       c=[cmap(lat_idx)], s=15, alpha=0.5,
                       label=f'$l_{{{lat_idx+1}}}$')

        ax.set_title(METHOD_LABELS.get(method, method), fontsize=14, fontweight='bold')
        ax.set_xlabel(f'PC1 ({pca.explained_variance_ratio_[0]:.1%})')
        ax.set_ylabel(f'PC2 ({pca.explained_variance_ratio_[1]:.1%})')
        if m_idx == n_methods - 1:
            ax.legend(bbox_to_anchor=(1.02, 1), loc='upper left',
                      fontsize=9, markerscale=1.5)

    plt.tight_layout()
    plt.savefig(output_path, dpi=300)
    plt.savefig(output_path.replace('.png', '.pdf'), dpi=300)
    plt.close()
    print(f'  Saved: {output_path}')


# ─────────────────────────────────────────────────────────────────────────────
# Main
# ─────────────────────────────────────────────────────────────────────────────

def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('--pmps_dir', type=str, required=True)
    parser.add_argument('--codi_dir', type=str, required=True)
    parser.add_argument('--simcot_dir', type=str, default=None)
    parser.add_argument('--output_dir', type=str, default='analysis/figures/collapse')
    args = parser.parse_args()

    os.makedirs(args.output_dir, exist_ok=True)

    methods = ['codi', 'pmps']
    dirs = {'codi': args.codi_dir, 'pmps': args.pmps_dir}
    if args.simcot_dir:
        methods.append('simcot')
        dirs['simcot'] = args.simcot_dir

    print('Loading results...')
    metrics_dict = {}
    hiddens_dict = {}
    for method in methods:
        m, h = load_results(dirs[method])
        metrics_dict[method] = m
        hiddens_dict[method] = h
        print(f'  {method}: {len(m)} samples, latent shape {h.shape}')

    print('\nGenerating figures...')

    plot_similarity_heatmaps(
        metrics_dict,
        os.path.join(args.output_dir, 'fig_similarity_heatmap.png'),
        methods=methods,
    )

    plot_similarity_distribution(
        metrics_dict,
        os.path.join(args.output_dir, 'fig_similarity_distribution.png'),
        methods=methods,
    )

    plot_singular_values(
        metrics_dict,
        os.path.join(args.output_dir, 'fig_singular_values.png'),
        methods=methods,
    )

    plot_effective_rank(
        metrics_dict,
        os.path.join(args.output_dir, 'fig_effective_rank.png'),
        methods=methods,
    )

    plot_tsne_visualization(
        hiddens_dict,
        os.path.join(args.output_dir, 'fig_tsne.png'),
        methods=methods,
    )

    plot_pca_visualization(
        hiddens_dict,
        os.path.join(args.output_dir, 'fig_pca.png'),
        methods=methods,
    )

    print(f'\nAll figures saved to: {args.output_dir}')


if __name__ == '__main__':
    main()