Blame for code/pmps/analysis/plot_distribution_analysis.py

"""
Information Distribution Analysis Visualization (PMPS vs Sim-CoT)

Generates publication-quality figures for information distribution uniformity:
  Fig 1: Information load per latent token (bar chart comparison)
  Fig 2: Matching weight heatmap (PMPS, case studies)
  Fig 3: Prototype utilization histogram
  Fig 4: Latent-CoT coverage heatmap

Usage:
    python analysis/plot_distribution_analysis.py \
        --pmps_dir analysis/results/pmps \
        --simcot_dir analysis/results/simcot \
        --output_dir analysis/figures/distribution
"""

import os
import json
import argparse
import numpy as np
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
from matplotlib.colors import LinearSegmentedColormap

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,
})

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

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


# ─────────────────────────────────────────────────────────────────────────────
# Fig 1: Information Load per Latent Token
# ─────────────────────────────────────────────────────────────────────────────

def plot_info_load_comparison(pmps_dist_metrics, simcot_info_loads, output_path,
                              num_latent=6):
    """
    Compare information load per latent token between PMPS and Sim-CoT.

    PMPS: uses matching weight sums from A_c2l
    Sim-CoT: uses uniform segmentation (N_c / N_l per latent, with remainder
             going to the last latent)
    """
    fig, axes = plt.subplots(1, 2, figsize=(11, 4.5))

    # --- Left: Per-latent info load bar chart ---
    ax = axes[0]

    # PMPS info loads
    pmps_loads = np.array([m['info_load_per_latent'] for m in pmps_dist_metrics])
    pmps_mean = pmps_loads.mean(axis=0)
    pmps_std = pmps_loads.std(axis=0)
    N_l = len(pmps_mean)

    # Sim-CoT info loads (if provided, else simulate rigid segmentation)
    if simcot_info_loads is not None:
        simcot_loads = np.array(simcot_info_loads)
        simcot_mean = simcot_loads.mean(axis=0)
        simcot_std = simcot_loads.std(axis=0)
    else:
        simcot_mean = None

    x = np.arange(N_l)
    w = 0.35

    bars_pmps = ax.bar(x - w/2, pmps_mean, w, yerr=pmps_std, capsize=4,
                       color=COLORS['pmps'], alpha=0.8, label=METHOD_LABELS['pmps'])

    if simcot_mean is not None:
        bars_simcot = ax.bar(x + w/2, simcot_mean, w, yerr=simcot_std, capsize=4,
                            color=COLORS['simcot'], alpha=0.8, label=METHOD_LABELS['simcot'])

    ideal_load = pmps_mean.sum() / N_l
    ax.axhline(y=ideal_load, color='gray', linestyle='--', alpha=0.6,
               label=f'Ideal (uniform = {ideal_load:.1f})')

    ax.set_xlabel('Latent Token Index')
    ax.set_ylabel('Information Load\n($\\sum_j A_{c2l}(j, i)$)')
    ax.set_title('(a) Information Load per Latent', fontsize=13)
    ax.set_xticks(x)
    ax.set_xticklabels([f'$l_{{{i+1}}}$' for i in range(N_l)])
    ax.legend(fontsize=9)

    # --- Right: CV (Coefficient of Variation) comparison ---
    ax2 = axes[1]

    pmps_cvs = [m['info_load_cv'] for m in pmps_dist_metrics]

    categories = [METHOD_LABELS['pmps']]
    means = [np.mean(pmps_cvs)]
    stds = [np.std(pmps_cvs)]
    bar_colors = [COLORS['pmps']]
    scatter_data = [pmps_cvs]

    if simcot_info_loads is not None:
        simcot_cvs = [np.std(loads) / (np.mean(loads) + 1e-8) for loads in simcot_info_loads]
        categories.append(METHOD_LABELS['simcot'])
        means.append(np.mean(simcot_cvs))
        stds.append(np.std(simcot_cvs))
        bar_colors.append(COLORS['simcot'])
        scatter_data.append(simcot_cvs)

    bars = ax2.bar(range(len(categories)), means, yerr=stds, capsize=5,
                   color=bar_colors, alpha=0.7, width=0.5)

    for i, (d, color) in enumerate(zip(scatter_data, bar_colors)):
        jitter = np.random.normal(0, 0.03, len(d))
        ax2.scatter(i + jitter, d, color=color, alpha=0.15, s=10, zorder=3)
        ax2.text(i, means[i] + stds[i] + 0.02, f'{means[i]:.3f}',
                ha='center', fontsize=11, fontweight='bold', color=color)

    ax2.set_xticks(range(len(categories)))
    ax2.set_xticklabels(categories, fontsize=12)
    ax2.set_ylabel('Coefficient of Variation\n(Lower = More Uniform)', fontsize=11)
    ax2.set_title('(b) Load Uniformity (CV)', fontsize=13)

    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 2: Matching Weight Heatmap (Case Studies)
# ─────────────────────────────────────────────────────────────────────────────

def plot_matching_weight_heatmaps(pmps_internals, output_path, num_cases=3):
    """
    Heatmap of A_c2l for individual samples.
    Shows which CoT tokens are assigned to which latent tokens.
    """
    n_cases = min(num_cases, len(pmps_internals))
    fig, axes = plt.subplots(1, n_cases, figsize=(5 * n_cases, 5))
    if n_cases == 1:
        axes = [axes]

    cmap = LinearSegmentedColormap.from_list(
        'weight', ['#FFFFFF', '#E8F8F5', '#1ABC9C', '#117864'], N=256
    )

    for idx in range(n_cases):
        sample = pmps_internals[idx]
        A_c2l = np.array(sample['A_c2l'])  # (N_c, N_l)
        N_c, N_l = A_c2l.shape

        ax = axes[idx]
        im = ax.imshow(A_c2l, cmap=cmap, aspect='auto', vmin=0)

        ax.set_xlabel('Latent Token Index')
        ax.set_ylabel('CoT Token Position')
        ax.set_title(f'Sample #{sample["sample_idx"]}', fontsize=12, fontweight='bold')
        ax.set_xticks(range(N_l))
        ax.set_xticklabels([f'$l_{{{i+1}}}$' for i in range(N_l)])

        if N_c > 30:
            tick_step = max(1, N_c // 10)
            ax.set_yticks(range(0, N_c, tick_step))

    fig.colorbar(im, ax=axes, label='Matching Weight $A_{c2l}$', shrink=0.8, pad=0.02)
    fig.suptitle('CoT → Latent Matching Weights (PMPS)', fontsize=14, y=1.02)

    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: Prototype Utilization
# ─────────────────────────────────────────────────────────────────────────────

def plot_prototype_utilization(pmps_dist_metrics, output_path):
    """
    Prototype utilization histogram + entropy comparison.
    """
    fig, axes = plt.subplots(1, 2, figsize=(11, 4.5))

    # Left: Average prototype usage (from latent assignments)
    ax = axes[0]
    all_proto_usage = np.array([m['proto_usage_latent'] for m in pmps_dist_metrics])
    avg_usage = all_proto_usage.mean(axis=0)
    std_usage = all_proto_usage.std(axis=0)
    K = len(avg_usage)
    uniform_level = 1.0 / K

    ax.bar(range(K), avg_usage, yerr=std_usage, capsize=2,
           color=COLORS['pmps'], alpha=0.7, edgecolor='white', linewidth=0.5)
    ax.axhline(y=uniform_level, color='red', linestyle='--', alpha=0.6,
               label=f'Uniform = {uniform_level:.4f}')
    ax.set_xlabel('Prototype Index')
    ax.set_ylabel('Average Assignment Weight')
    ax.set_title('(a) Prototype Utilization (Latent Side)', fontsize=13)
    ax.legend(fontsize=9)

    if K > 16:
        ax.set_xticks(range(0, K, max(1, K // 8)))

    # Right: Entropy comparison
    ax2 = axes[1]
    proto_ents = [m['proto_entropy_latent'] for m in pmps_dist_metrics]
    max_ent = pmps_dist_metrics[0]['max_entropy']

    ax2.hist(proto_ents, bins=30, color=COLORS['pmps'], alpha=0.7,
             edgecolor='white', linewidth=0.5, density=True)
    ax2.axvline(x=max_ent, color='red', linestyle='--', linewidth=2, alpha=0.7,
                label=f'Max Entropy = $\\log K$ = {max_ent:.2f}')
    ax2.axvline(x=np.mean(proto_ents), color=COLORS['pmps'], linestyle='-',
                linewidth=2, label=f'Mean = {np.mean(proto_ents):.2f}')

    ax2.set_xlabel('Prototype Assignment Entropy')
    ax2.set_ylabel('Density')
    ax2.set_title('(b) Prototype Entropy Distribution', fontsize=13)
    ax2.legend(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 4: Latent-CoT Coverage Comparison
# ─────────────────────────────────────────────────────────────────────────────

def plot_coverage_comparison(pmps_dist_metrics, simcot_info_loads, output_path):
    """
    Compare how CoT information is distributed across latents.

    For PMPS: use matching weight entropy per CoT token (higher = distributed more evenly)
    For Sim-CoT: compute rigid-segment coverage
    """
    fig, ax = plt.subplots(figsize=(7, 4.5))

    # PMPS matching entropies
    pmps_match_ents = []
    for m in pmps_dist_metrics:
        pmps_match_ents.extend(m['match_entropies'])

    data_list = [pmps_match_ents]
    labels = [METHOD_LABELS['pmps']]
    colors = [COLORS['pmps']]

    # For Sim-CoT: approximate with delta function (each CoT token maps to exactly 1 latent)
    # entropy = 0 for rigid one-to-one mapping
    if simcot_info_loads is not None:
        simcot_ents = [0.0] * len(pmps_match_ents)
        data_list.append(simcot_ents)
        labels.append(METHOD_LABELS['simcot'])
        colors.append(COLORS['simcot'])

    parts = ax.violinplot(data_list, positions=range(len(labels)),
                          showmedians=True, showextrema=False)
    for pc, color in zip(parts['bodies'], colors):
        pc.set_facecolor(color)
        pc.set_alpha(0.4)
    parts['cmedians'].set_color('black')
    parts['cmedians'].set_linewidth(2)

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

    ax.set_xticks(range(len(labels)))
    ax.set_xticklabels(labels, fontsize=12)
    ax.set_ylabel('Per-CoT-Token Matching Entropy\n(Higher = More Distributed)', fontsize=11)
    ax.set_title('Matching Weight Distribution\n(How Diversely Each CoT Token Maps to Latents)', fontsize=13)

    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('--simcot_dir', type=str, default=None)
    parser.add_argument('--output_dir', type=str, default='analysis/figures/distribution')
    args = parser.parse_args()

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

    # Load PMPS distribution metrics
    print('Loading PMPS distribution metrics...')
    pmps_dist_path = os.path.join(args.pmps_dir, 'distribution_metrics.json')
    pmps_pmps_path = os.path.join(args.pmps_dir, 'pmps_internals.json')

    with open(pmps_dist_path, 'r') as f:
        pmps_dist_metrics = json.load(f)
    print(f'  PMPS: {len(pmps_dist_metrics)} samples with distribution metrics')

    pmps_internals = None
    if os.path.exists(pmps_pmps_path):
        with open(pmps_pmps_path, 'r') as f:
            pmps_internals = json.load(f)
        print(f'  PMPS: {len(pmps_internals)} samples with PMPS internals')

    # Load Sim-CoT info loads (if available)
    simcot_info_loads = None
    if args.simcot_dir:
        simcot_dist_path = os.path.join(args.simcot_dir, 'distribution_metrics.json')
        if os.path.exists(simcot_dist_path):
            with open(simcot_dist_path, 'r') as f:
                simcot_dist = json.load(f)
            simcot_info_loads = [m['info_load_per_latent'] for m in simcot_dist]
            print(f'  Sim-CoT: {len(simcot_info_loads)} samples')
        else:
            print('  Sim-CoT: no distribution_metrics.json found, will simulate rigid segmentation')

    print('\nGenerating figures...')

    plot_info_load_comparison(
        pmps_dist_metrics,
        simcot_info_loads,
        os.path.join(args.output_dir, 'fig_info_load.png'),
    )

    if pmps_internals:
        plot_matching_weight_heatmaps(
            pmps_internals,
            os.path.join(args.output_dir, 'fig_matching_weights.png'),
            num_cases=3,
        )

    plot_prototype_utilization(
        pmps_dist_metrics,
        os.path.join(args.output_dir, 'fig_prototype_utilization.png'),
    )

    plot_coverage_comparison(
        pmps_dist_metrics,
        simcot_info_loads,
        os.path.join(args.output_dir, 'fig_coverage_comparison.png'),
    )

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


if __name__ == '__main__':
    main()