"""
Paper-Quality Combined Analysis Figures for PMPS

Generates the main analysis figures for the AAAI paper:
  Fig A: Three-way comparison (CODI vs Sim-CoT vs PMPS) — the hero figure
  Fig B: Matching weight case study with CoT annotation
  Fig C: Training dynamics (collapse metrics over training steps)

Usage:
    python analysis/plot_paper_figures.py \
        --pmps_dir analysis/results/pmps \
        --codi_dir analysis/results/codi \
        --simcot_dir analysis/results/simcot \
        --output_dir analysis/paper_figures
"""

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 matplotlib.patches import FancyBboxPatch

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

COLORS = {
    'codi': '#E74C3C',
    'pmps': '#2E86C1',
    'simcot': '#27AE60',
}
METHOD_LABELS = {
    'codi': 'CODI',
    'pmps': 'PMPS (Ours)',
    'simcot': 'Sim-CoT',
}


def load_metrics(result_dir):
    with open(os.path.join(result_dir, 'metrics.json'), 'r') as f:
        return json.load(f)


def load_dist_metrics(result_dir):
    path = os.path.join(result_dir, 'distribution_metrics.json')
    if os.path.exists(path):
        with open(path, 'r') as f:
            return json.load(f)
    return None


def load_latent_hiddens(result_dir):
    return np.load(os.path.join(result_dir, 'latent_hiddens.npy'))


# ─────────────────────────────────────────────────────────────────────────────
# Figure A: Three-way Comparison — The Hero Figure
#
#   (a) Cosine Similarity Heatmaps (3 side-by-side)
#   (b) Effective Rank bar chart
#   (c) Information Load per Latent (PMPS vs Sim-CoT)
# ─────────────────────────────────────────────────────────────────────────────

def plot_hero_figure(metrics_dict, dist_metrics_dict, output_path,
                     methods=['codi', 'simcot', 'pmps']):
    """The main analysis figure for the paper."""

    available_methods = [m for m in methods if m in metrics_dict]

    fig = plt.figure(figsize=(14, 4.8))
    gs = gridspec.GridSpec(1, 3, figure=fig, wspace=0.35,
                           width_ratios=[2.5, 1, 1.5])

    # ── (a) Cosine Similarity Heatmaps ──────────────────────────────────────
    gs_heatmaps = gs[0].subgridspec(1, len(available_methods), wspace=0.15)

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

    for m_idx, method in enumerate(available_methods):
        ax = fig.add_subplot(gs_heatmaps[m_idx])
        metrics = metrics_dict[method]
        sim_matrices = [np.array(m['cos_sim_matrix']) for m in metrics]
        avg_matrix = np.mean(sim_matrices, axis=0)
        N_l = avg_matrix.shape[0]

        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.65 else 'black'
                ax.text(j, i, f'{avg_matrix[i, j]:.2f}',
                        ha='center', va='center', fontsize=7, color=color)

        ax.set_title(METHOD_LABELS.get(method, method), fontsize=11, fontweight='bold',
                     color=COLORS.get(method, 'black'))
        ax.set_xticks(range(N_l))
        ax.set_yticks(range(N_l))
        ax.set_xticklabels([f'{i+1}' for i in range(N_l)], fontsize=8)
        if m_idx == 0:
            ax.set_yticklabels([f'{i+1}' for i in range(N_l)], fontsize=8)
            ax.set_ylabel('Latent Index', fontsize=10)
        else:
            ax.set_yticklabels([])
        ax.set_xlabel('Latent Index', fontsize=10)

        avg_offdiag = np.mean([m['avg_off_diag_sim'] for m in metrics])
        ax.text(0.5, -0.22, f'Avg off-diag: {avg_offdiag:.3f}',
                transform=ax.transAxes, ha='center', fontsize=8,
                color=COLORS.get(method, 'black'), fontweight='bold')

    # ── (b) Effective Rank ───────────────────────────────────────────────────
    ax_rank = fig.add_subplot(gs[1])

    for m_idx, method in enumerate(available_methods):
        ranks = [m['effective_rank'] for m in metrics_dict[method]]
        mean_r = np.mean(ranks)
        std_r = np.std(ranks)
        color = COLORS.get(method, '#888')

        ax_rank.bar(m_idx, mean_r, width=0.6, yerr=std_r, capsize=4,
                    color=color, alpha=0.75)
        ax_rank.text(m_idx, mean_r + std_r + 0.08, f'{mean_r:.2f}',
                     ha='center', fontsize=9, fontweight='bold', color=color)

    N_l = len(metrics_dict[available_methods[0]][0]['cos_sim_matrix'])
    ax_rank.axhline(y=N_l, color='gray', linestyle='--', alpha=0.4)
    ax_rank.text(len(available_methods) - 0.3, N_l + 0.05, f'Max={N_l}',
                 color='gray', fontsize=8)
    ax_rank.axhline(y=1.0, color='#E74C3C', linestyle=':', alpha=0.4)
    ax_rank.text(len(available_methods) - 0.3, 1.1, 'Collapse', color='#E74C3C', fontsize=8)

    ax_rank.set_xticks(range(len(available_methods)))
    ax_rank.set_xticklabels([METHOD_LABELS.get(m, m) for m in available_methods],
                            fontsize=9, rotation=15)
    ax_rank.set_ylabel('Effective Rank')
    ax_rank.set_title('(b) Effective Rank', fontsize=11)
    ax_rank.set_ylim(0, N_l + 1.5)

    # ── (c) Information Load Uniformity ──────────────────────────────────────
    ax_load = fig.add_subplot(gs[2])

    load_methods = [m for m in available_methods if m in dist_metrics_dict and dist_metrics_dict[m]]

    if load_methods:
        bar_width = 0.7 / len(load_methods)
        for lm_idx, method in enumerate(load_methods):
            dist_m = dist_metrics_dict[method]
            loads = np.array([m['info_load_per_latent'] for m in dist_m])
            mean_loads = loads.mean(axis=0)
            std_loads = loads.std(axis=0)
            N_l = len(mean_loads)
            x = np.arange(N_l) + lm_idx * bar_width
            color = COLORS.get(method, '#888')
            ax_load.bar(x, mean_loads, bar_width, yerr=std_loads, capsize=2,
                        color=color, alpha=0.75,
                        label=METHOD_LABELS.get(method, method))

        ideal = mean_loads.sum() / N_l
        ax_load.axhline(y=ideal, color='gray', linestyle='--', alpha=0.5,
                        label=f'Ideal = {ideal:.1f}')
        ax_load.set_xticks(np.arange(N_l) + bar_width * (len(load_methods) - 1) / 2)
        ax_load.set_xticklabels([f'$l_{{{i+1}}}$' for i in range(N_l)], fontsize=9)
        ax_load.set_xlabel('Latent Index')
        ax_load.set_ylabel('Information Load')
        ax_load.set_title('(c) Info Distribution', fontsize=11)
        ax_load.legend(fontsize=8, loc='upper right')
    else:
        ax_load.text(0.5, 0.5, 'No distribution\ndata available',
                     transform=ax_load.transAxes, ha='center', va='center',
                     fontsize=11, color='gray')
        ax_load.set_title('(c) Info Distribution', fontsize=11)

    # Add overall title
    fig.text(0.18, 1.02, '(a) Latent Pairwise Cosine Similarity',
             fontsize=11, ha='center', transform=fig.transFigure)

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


# ─────────────────────────────────────────────────────────────────────────────
# Figure B: Matching Weight Case Study
# ─────────────────────────────────────────────────────────────────────────────

def plot_matching_case_study(pmps_internals_path, output_path, num_cases=2):
    """Detailed matching weight visualization for individual samples."""
    with open(pmps_internals_path, 'r') as f:
        pmps_internals = json.load(f)

    n_cases = min(num_cases, len(pmps_internals))
    fig, axes = plt.subplots(1, n_cases, figsize=(5.5 * n_cases, 5))
    if n_cases == 1:
        axes = [axes]

    cmap = LinearSegmentedColormap.from_list(
        'weight', ['#FDFEFE', '#D5F5E3', '#27AE60', '#1E8449', '#145A32'], N=256
    )

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

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

        ax.set_xlabel('Latent Token')
        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}}}$\n({info_load[i]:.1f})' for i in range(N_l)],
                           fontsize=9)

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

    fig.colorbar(im, ax=axes, label='$A_{c2l}$ Weight', shrink=0.8, pad=0.03)
    fig.suptitle('PMPS Matching Weights: CoT Token → Latent Token\n'
                 '(Numbers below latent labels show information load)',
                 fontsize=12, y=1.05)

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


# ─────────────────────────────────────────────────────────────────────────────
# Figure C: Singular Value Comparison (compact version)
# ─────────────────────────────────────────────────────────────────────────────

def plot_sv_compact(metrics_dict, output_path, methods=None):
    """Compact singular value comparison for supplementary material."""
    if methods is None:
        methods = list(metrics_dict.keys())

    fig, axes = plt.subplots(1, 2, figsize=(10, 4))

    # Left: Normalized singular values
    ax = axes[0]
    for method in methods:
        sv_norms = np.array([m['singular_values_normalized'] for m in metrics_dict[method]])
        mean_sv = sv_norms.mean(axis=0)
        std_sv = sv_norms.std(axis=0)
        N_l = len(mean_sv)
        x = np.arange(1, N_l + 1)
        color = COLORS.get(method, '#888')
        ax.plot(x, mean_sv, 'o-', color=color, linewidth=2, markersize=6,
                label=METHOD_LABELS.get(method, method))
        ax.fill_between(x, mean_sv - std_sv, mean_sv + std_sv,
                         alpha=0.15, color=color)

    ax.set_xlabel('Singular Value Index')
    ax.set_ylabel('Normalized Magnitude')
    ax.set_title('(a) Singular Value Spectrum')
    ax.legend()
    ax.set_xticks(range(1, N_l + 1))

    # Right: Effective rank histogram
    ax2 = axes[1]
    for method in methods:
        ranks = [m['effective_rank'] for m in metrics_dict[method]]
        color = COLORS.get(method, '#888')
        ax2.hist(ranks, bins=20, color=color, alpha=0.4, edgecolor=color,
                 linewidth=1.2, label=METHOD_LABELS.get(method, method), density=True)
        ax2.axvline(x=np.mean(ranks), color=color, linestyle='--', linewidth=2)

    ax2.set_xlabel('Effective Rank')
    ax2.set_ylabel('Density')
    ax2.set_title('(b) Effective Rank Distribution')
    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}')


# ─────────────────────────────────────────────────────────────────────────────
# Figure D: Prototype Analysis (supplementary)
# ─────────────────────────────────────────────────────────────────────────────

def plot_prototype_analysis(pmps_dist_metrics, output_path):
    """Prototype utilization and entropy analysis."""
    fig, axes = plt.subplots(1, 3, figsize=(14, 4))

    # (a) Prototype utilization (latent side)
    ax = axes[0]
    all_usage_lat = np.array([m['proto_usage_latent'] for m in pmps_dist_metrics])
    avg_lat = all_usage_lat.mean(axis=0)
    K = len(avg_lat)
    ax.bar(range(K), avg_lat, color=COLORS['pmps'], alpha=0.7, edgecolor='white')
    ax.axhline(y=1.0/K, color='red', linestyle='--', alpha=0.6, label='Uniform')
    ax.set_xlabel('Prototype Index')
    ax.set_ylabel('Avg Weight')
    ax.set_title('(a) Proto Usage (Latent)')
    ax.legend(fontsize=8)

    # (b) Prototype utilization (CoT side)
    ax2 = axes[1]
    all_usage_cot = np.array([m['proto_usage_cot'] for m in pmps_dist_metrics])
    avg_cot = all_usage_cot.mean(axis=0)
    ax2.bar(range(K), avg_cot, color=COLORS['simcot'], alpha=0.7, edgecolor='white')
    ax2.axhline(y=1.0/K, color='red', linestyle='--', alpha=0.6, label='Uniform')
    ax2.set_xlabel('Prototype Index')
    ax2.set_ylabel('Avg Weight')
    ax2.set_title('(b) Proto Usage (CoT)')
    ax2.legend(fontsize=8)

    # (c) Entropy scatter
    ax3 = axes[2]
    lat_ents = [m['proto_entropy_latent'] for m in pmps_dist_metrics]
    cot_ents = [m['proto_entropy_cot'] for m in pmps_dist_metrics]
    max_ent = pmps_dist_metrics[0]['max_entropy']

    ax3.scatter(lat_ents, cot_ents, alpha=0.3, s=15, color=COLORS['pmps'])
    ax3.plot([0, max_ent], [0, max_ent], 'k--', alpha=0.3, label='y=x')
    ax3.axhline(y=max_ent, color='red', linestyle=':', alpha=0.3)
    ax3.axvline(x=max_ent, color='red', linestyle=':', alpha=0.3)
    ax3.set_xlabel('Latent Proto Entropy')
    ax3.set_ylabel('CoT Proto Entropy')
    ax3.set_title(f'(c) Entropy (max={max_ent:.2f})')
    ax3.legend(fontsize=8)
    ax3.set_aspect('equal')

    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/paper_figures')
    args = parser.parse_args()

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

    # Load all results
    print('Loading results...')
    methods = ['codi', 'pmps']
    dirs = {'codi': args.codi_dir, 'pmps': args.pmps_dir}
    if args.simcot_dir:
        methods.append('simcot')
        dirs['simcot'] = args.simcot_dir

    # Reorder: codi, simcot, pmps
    method_order = ['codi']
    if 'simcot' in methods:
        method_order.append('simcot')
    method_order.append('pmps')

    metrics_dict = {}
    dist_metrics_dict = {}

    for method in method_order:
        metrics_dict[method] = load_metrics(dirs[method])
        dist_metrics_dict[method] = load_dist_metrics(dirs[method])
        print(f'  {method}: {len(metrics_dict[method])} samples')

    # Generate figures
    print('\nGenerating paper figures...')

    # Hero figure (3-way comparison)
    plot_hero_figure(
        metrics_dict, dist_metrics_dict,
        os.path.join(args.output_dir, 'fig_analysis_main.png'),
        methods=method_order,
    )

    # Singular value comparison
    plot_sv_compact(
        metrics_dict,
        os.path.join(args.output_dir, 'fig_singular_values.png'),
        methods=method_order,
    )

    # Matching weight case study (PMPS only)
    pmps_pmps_path = os.path.join(args.pmps_dir, 'pmps_internals.json')
    if os.path.exists(pmps_pmps_path):
        plot_matching_case_study(
            pmps_pmps_path,
            os.path.join(args.output_dir, 'fig_matching_case_study.png'),
            num_cases=3,
        )

    # Prototype analysis (PMPS only)
    if dist_metrics_dict.get('pmps'):
        plot_prototype_analysis(
            dist_metrics_dict['pmps'],
            os.path.join(args.output_dir, 'fig_prototype_analysis.png'),
        )

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


if __name__ == '__main__':
    main()