#!/usr/bin/env python3
"""
Generate Figure 0: Universal Model Inadequacy on Hard Datasets

Shows RA scores for all models across 4 hardest datasets, revealing universal
failure in relational reasoning on challenging semantic categories.

Usage:
    python generate_figure0_hard_datasets_heatmap.py \
        --input cli_batch_benchmark_ra-2025-12-31-all-models-all-datasets-EN-ZH.csv \
        --output neurips_figure0_heatmap_hard_EN.pdf \
        --language EN \
        [--dpi 300]
"""

import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns
import click
from pathlib import Path


# Hard datasets for Figure 0 (confirmed from difficulty analysis)
# v2.5: Updated dataset names with -v2.5 suffix, removed NeurIPS-07-math-numbers (v1 only)
HARD_DATASETS = [
    'NeurIPS-07-sequential-mathematical-v2.5',
    'NeurIPS-11-cause-effect-v2.5',
    'NeurIPS-09-part-whole-v2.5',
]

DATASET_LABELS = [
    'Sequential-Mathematical',
    'Cause-Effect',
    'Part-Whole',
]


def create_heatmap_hard_datasets(
    data_path: str,
    output_path: str,
    language: str,  # 'EN', 'ZH', 'ES', or 'cross'
    figsize: tuple = (10, 10),
    dpi: int = 300,
    verbose: bool = True
):
    """
    Create heatmap showing RA scores on hardest datasets (universal inadequacy).

    Args:
        data_path: Path to RA raw CSV
        output_path: Path to save figure
        language: 'EN', 'ZH', 'ES', or 'cross'
        figsize: Figure size (width, height) in inches
        dpi: Resolution for output
        verbose: Print progress messages
    """
    if verbose:
        lang_name = 'English' if language == 'EN' else 'Chinese' if language == 'ZH' else 'Spanish' if language == 'ES' else 'Cross-lingual'
        click.echo(f"📊 Generating Figure 0 ({lang_name}): Hard Datasets Heatmap")
        click.echo(f"   Reading data: {data_path}")

    # Read raw RA data
    df = pd.read_csv(data_path)

    # Select RA column based on language
    ra_col = f'cos_{language}'

    if verbose:
        click.echo(f"   Using RA column: {ra_col}")
        click.echo(f"   Hard datasets: {', '.join(HARD_DATASETS)}")

    # Get unique models
    models = df['model'].unique().tolist()

    # Create matrix: rows=models, columns=hard datasets
    score_matrix = []
    model_labels = []

    for model in models:
        model_name = model.replace(' (OpenRouter)', '').replace(' (Ollama)', '')
        model_labels.append(model_name)

        row_scores = []
        for dataset in HARD_DATASETS:
            # Find RA score for this model-dataset pair
            match = df[(df['model'] == model) & (df['dataset'] == dataset)]
            if len(match) > 0:
                score = match.iloc[0][ra_col]
                row_scores.append(score)
            else:
                row_scores.append(np.nan)
        score_matrix.append(row_scores)

    # Create numpy array
    data = np.array(score_matrix)

    # Calculate average RA for tier assignment
    avg_ra = np.nanmean(data, axis=1)
    sorted_indices = np.argsort(avg_ra)[::-1]  # Sort descending

    # Reorder data and models by average RA
    data = data[sorted_indices]
    model_labels = [model_labels[i] for i in sorted_indices]
    avg_ra = avg_ra[sorted_indices]

    if verbose:
        click.echo(f"   Heatmap dimensions: {len(model_labels)} models × {len(DATASET_LABELS)} datasets")
        click.echo(f"   RA score range: {np.nanmin(data):.4f} to {np.nanmax(data):.4f}")
        click.echo(f"   Mean RA across all cells: {np.nanmean(data):.4f}")

    # Set publication style
    plt.style.use('seaborn-v0_8-white')
    sns.set_context("paper", font_scale=1.1)

    # Create figure
    fig, ax = plt.subplots(figsize=figsize, dpi=dpi)

    # Create heatmap (same color scheme as other figures for consistency)
    im = ax.imshow(data, cmap='RdYlGn', aspect='auto', vmin=0, vmax=0.8)

    # Set ticks and labels
    ax.set_xticks(np.arange(len(DATASET_LABELS)))
    ax.set_yticks(np.arange(len(model_labels)))
    ax.set_xticklabels(DATASET_LABELS, rotation=45, ha='right', fontsize=11)
    ax.set_yticklabels(model_labels, rotation=45, ha='right', fontsize=10)

    # Disable grid
    ax.grid(False)

    # Add colorbar
    cbar = plt.colorbar(im, ax=ax, fraction=0.046, pad=0.04)
    cbar.set_label('Relational Affinity (RA)', rotation=270, labelpad=20, fontsize=12, fontweight='bold')

    # Annotate cells with values
    for i in range(len(model_labels)):
        for j in range(len(DATASET_LABELS)):
            val = data[i, j]
            if not np.isnan(val):
                text = ax.text(j, i, f'{val:.3f}',
                             ha="center", va="center", color='black', fontsize=13, fontweight='bold')

    # Add tier separator lines (expect mostly Tier 3 - red)
    tier1_end = np.sum(avg_ra > 0.50)
    tier2_end = tier1_end + np.sum((avg_ra >= 0.30) & (avg_ra <= 0.50))

    if tier1_end > 0 and tier1_end < len(model_labels):
        ax.axhline(y=tier1_end - 0.5, color='black', linestyle='-', linewidth=2.5)
        ax.text(-0.5, tier1_end / 2 - 0.5, 'Tier 1', rotation=90, va='center', ha='center',
                fontsize=11, fontweight='bold', color='green')

    if tier2_end > tier1_end and tier2_end < len(model_labels):
        ax.axhline(y=tier2_end - 0.5, color='black', linestyle='-', linewidth=2.5)
        ax.text(-0.5, (tier1_end + tier2_end) / 2 - 0.5, 'Tier 2', rotation=90, va='center', ha='center',
                fontsize=11, fontweight='bold', color='orange')

    if tier2_end < len(model_labels):
        ax.text(-0.5, (tier2_end + len(model_labels)) / 2 - 0.5, 'Tier 3', rotation=90, va='center', ha='center',
                fontsize=11, fontweight='bold', color='red')

    # Labels and title
    lang_name = 'English' if language == 'EN' else 'Chinese' if language == 'ZH' else 'Spanish' if language == 'ES' else 'Cross-lingual'
    ax.set_xlabel('Semantic Category', fontsize=13, fontweight='bold')
    ax.set_ylabel('Embedding Model (Ranked by Average RA)', fontsize=13, fontweight='bold')
    ax.set_title(
        f'Figure 1: Universal Model Inadequacy on Challenging Datasets ({lang_name})',
        fontsize=15,
        fontweight='bold',
        pad=20
    )

    # Tight layout
    plt.tight_layout()

    # Save figure
    output_path = Path(output_path)
    plt.savefig(output_path, dpi=dpi, bbox_inches='tight')

    if verbose:
        click.echo(f"\n   ✅ Figure saved: {output_path}")
        click.echo(f"   Resolution: {dpi} DPI")
        click.echo(f"   Size: {figsize[0]}\" × {figsize[1]}\"")

    # Also save as PNG
    png_path = output_path.with_suffix('.png')
    plt.savefig(png_path, dpi=150, bbox_inches='tight')

    if verbose:
        click.echo(f"   ✅ Preview saved: {png_path}")

    plt.close()

    # Print summary statistics
    if verbose:
        click.echo(f"\n📊 Universal Inadequacy Statistics ({lang_name}):")
        click.echo(f"   Overall mean RA: {np.nanmean(data):.4f}")
        click.echo(f"   Overall std RA:  {np.nanstd(data):.4f}")

        # Find best and worst datasets
        dataset_means = np.nanmean(data, axis=0)
        best_ds_idx = np.nanargmax(dataset_means)
        worst_ds_idx = np.nanargmin(dataset_means)
        click.echo(f"\n   Least-hard dataset:  {DATASET_LABELS[best_ds_idx]} (mean RA={dataset_means[best_ds_idx]:.4f})")
        click.echo(f"   Hardest dataset:     {DATASET_LABELS[worst_ds_idx]} (mean RA={dataset_means[worst_ds_idx]:.4f})")

        # Print tier counts (expect 0% Tier 1)
        tier1_count = np.sum(avg_ra > 0.50)
        tier2_count = np.sum((avg_ra >= 0.30) & (avg_ra <= 0.50))
        tier3_count = np.sum(avg_ra < 0.30)

        click.echo(f"\n   Tier 1 (RA > 0.50): {tier1_count} models ({tier1_count/len(model_labels)*100:.1f}%)")
        click.echo(f"   Tier 2 (0.30 ≤ RA ≤ 0.50): {tier2_count} models ({tier2_count/len(model_labels)*100:.1f}%)")
        click.echo(f"   Tier 3 (RA < 0.30): {tier3_count} models ({tier3_count/len(model_labels)*100:.1f}%)")

        if tier1_count == 0:
            click.echo("\n   ⚠️  CRITICAL: NO models achieve Tier 1 performance on hard datasets!")
            click.echo("   This confirms universal inadequacy in relational reasoning.")


@click.command()
@click.option(
    '--input',
    '-i',
    required=True,
    type=click.Path(exists=True, dir_okay=False),
    help='Path to RA raw CSV file'
)
@click.option(
    '--output',
    '-o',
    required=True,
    type=click.Path(dir_okay=False),
    help='Path to save output figure (PDF recommended)'
)
@click.option(
    '--language',
    '-l',
    required=True,
    type=click.Choice(['EN', 'ZH', 'ES', 'cross'], case_sensitive=False),
    help='Language to analyze: EN (English), ZH (Chinese), ES (Spanish), or cross (cross-lingual)'
)
@click.option(
    '--dpi',
    default=300,
    type=int,
    help='Resolution for output figure (default: 300)'
)
@click.option(
    '--figsize',
    default=(10, 10),
    type=(float, float),
    help='Figure size in inches (width, height) (default: 10 10)'
)
@click.option(
    '--verbose/--quiet',
    '-v/-q',
    default=True,
    help='Print progress messages'
)
def main(
    input: str,
    output: str,
    language: str,
    dpi: int,
    figsize: tuple,
    verbose: bool
):
    """
    Generate Figure 0: Hard datasets performance heatmap (universal inadequacy).

    Shows RA scores for all models on the 4 hardest datasets, revealing that
    NO model adequately captures relational logic on challenging categories.

    Examples:

        # English version
        python generate_figure0_hard_datasets_heatmap.py \\
            --input cli_batch_benchmark_ra-2025-12-31-all-models-all-datasets-EN-ZH.csv \\
            --output neurips_figure0_heatmap_hard_EN.pdf \\
            --language EN

        # Chinese version
        python generate_figure0_hard_datasets_heatmap.py \\
            --input cli_batch_benchmark_ra-2025-12-31-all-models-all-datasets-EN-ZH.csv \\
            --output neurips_figure0_heatmap_hard_ZH.pdf \\
            --language ZH

        # Spanish version
        python generate_figure0_hard_datasets_heatmap.py \\
            --input cli_batch_benchmark_ra-2025-12-31-all-models-all-datasets-EN-ES.csv \\
            --output neurips_figure0_heatmap_hard_ES.pdf \\
            --language ES
    """
    if verbose:
        lang_name = 'English' if language == 'EN' else 'Chinese' if language == 'ZH' else 'Spanish' if language == 'ES' else 'Cross-lingual'
        click.echo("\n" + "=" * 80)
        click.echo(f"FIGURE 0 GENERATION: Hard Datasets Heatmap ({lang_name})")
        click.echo("=" * 80 + "\n")

    create_heatmap_hard_datasets(
        data_path=input,
        output_path=output,
        language=language,
        figsize=figsize,
        dpi=dpi,
        verbose=verbose
    )

    if verbose:
        click.echo("\n" + "=" * 80)
        click.echo("✅ Figure 0 generation complete!")
        click.echo(f"   Main figure: {output}")
        click.echo(f"   Preview: {Path(output).with_suffix('.png')}")
        click.echo("=" * 80 + "\n")


if __name__ == '__main__':
    main()
