generate_comparison_charts.py

"""
Generate comparison charts for Financial Entity Extraction experiments.

Compares:
- OpenAI Baseline (direct API)
- DSPy Baseline
- DSPy + GEPA (optimized)
- DSPy + BAML Baseline
- DSPy + BAML + GEPA (optimized)
"""

import matplotlib.pyplot as plt
import numpy as np
import pandas as pd

# Data from all_stats.txt
# Overall metrics - ordered: Baselines together, then GEPA versions together
results_data = {
    "Method": [
        "OpenAI\nBaseline",
        "DSPy\nBaseline",
        "DSPy +\nBAML",
        "DSPy +\nGEPA",
        "DSPy + BAML\n+ GEPA"
    ],
    "Exact Match": [0.3207, 0.397875, 0.427391, 0.538371, 0.544274],
    "Mean Field": [0.8434, 0.870552, 0.878563, 0.916428, 0.916175]
}

# Per-field accuracy data
fields = ["Company", "Date", "Location", "Money", "Person", "Product", "Quantity"]

per_field_data = {
    "Field": fields,
    "OpenAI Baseline": [0.7369, 0.7752, 0.8975, 0.9372, 0.9825, 0.6915, 0.8829],
    "DSPy Baseline": [0.7314, 0.8288, 0.9327, 0.9587, 0.9911, 0.7473, 0.9038],
    "DSPy + GEPA": [0.7922, 0.9179, 0.9498, 0.9847, 0.9917, 0.8672, 0.9115],
    "DSPy + BAML": [0.7538, 0.8318, 0.9292, 0.9599, 0.9888, 0.7816, 0.9050],
    "DSPy + BAML + GEPA": [0.7834, 0.9185, 0.9534, 0.9811, 0.9923, 0.8619, 0.9227]
}

def create_overall_comparison_chart():
    """Create bar chart comparing overall metrics across all methods."""
    fig, axes = plt.subplots(1, 2, figsize=(16, 6))

    methods = results_data["Method"]
    x = np.arange(len(methods))

    # Color scheme - baselines in blues/purples, GEPA optimized in greens
    colors = ['#E74C3C', '#3498DB', '#9B59B6', '#2ECC71', '#1ABC9C']

    # Exact Match Accuracy
    ax1 = axes[0]
    bars1 = ax1.bar(x, results_data["Exact Match"], color=colors, edgecolor='black', linewidth=1.2)
    ax1.set_ylabel("Accuracy", fontsize=12)
    ax1.set_title("Exact Match Accuracy", fontsize=14, fontweight='bold')
    ax1.set_xticks(x)
    ax1.set_xticklabels(methods, fontsize=10)
    ax1.set_ylim(0, 0.7)
    ax1.axhline(y=0.490, color='gray', linestyle='--', alpha=0.5, linewidth=1)
    ax1.grid(axis='y', alpha=0.3)

    for bar, val in zip(bars1, results_data["Exact Match"]):
        ax1.text(bar.get_x() + bar.get_width()/2, bar.get_height() + 0.015,
                 f'{val:.1%}', ha='center', va='bottom', fontsize=11, fontweight='bold')

    # Mean Field Accuracy
    ax2 = axes[1]
    bars2 = ax2.bar(x, results_data["Mean Field"], color=colors, edgecolor='black', linewidth=1.2)
    ax2.set_ylabel("Accuracy", fontsize=12)
    ax2.set_title("Mean Field Accuracy", fontsize=14, fontweight='bold')
    ax2.set_xticks(x)
    ax2.set_xticklabels(methods, fontsize=10)
    ax2.set_ylim(0.8, 1.0)
    ax2.axhline(y=0.8825, color='gray', linestyle='--', alpha=0.5, linewidth=1)
    ax2.grid(axis='y', alpha=0.3)

    for bar, val in zip(bars2, results_data["Mean Field"]):
        ax2.text(bar.get_x() + bar.get_width()/2, bar.get_height() + 0.005,
                 f'{val:.1%}', ha='center', va='bottom', fontsize=11, fontweight='bold')

    plt.suptitle("Financial Entity Extraction: Method Comparison", fontsize=16, fontweight='bold', y=1.02)
    plt.tight_layout()
    plt.savefig("comparison_overall.png", dpi=150, bbox_inches='tight', facecolor='white')
    plt.show()
    print("Saved: comparison_overall.png")


def create_per_field_comparison_chart():
    """Create grouped bar chart comparing per-field accuracy."""
    fig, ax = plt.subplots(figsize=(16, 7))

    x = np.arange(len(fields))
    width = 0.15

    # All methods including OpenAI Baseline - baselines first, then GEPA versions together
    methods = ["OpenAI Baseline", "DSPy Baseline", "DSPy + BAML", "DSPy + GEPA", "DSPy + BAML + GEPA"]
    colors = ['#E74C3C', '#3498DB', '#9B59B6', '#2ECC71', '#1ABC9C']

    for i, (method, color) in enumerate(zip(methods, colors)):
        offset = (i - 2) * width
        bars = ax.bar(x + offset, per_field_data[method], width,
                     label=method, color=color, edgecolor='black', linewidth=0.8)

    ax.set_ylabel('Accuracy', fontsize=12)
    ax.set_title('Per-Field Accuracy Comparison', fontsize=14, fontweight='bold')
    ax.set_xticks(x)
    ax.set_xticklabels(fields, fontsize=11)
    ax.set_ylim(0.65, 1.02)
    ax.grid(axis='y', alpha=0.3)

    # Place legend outside the plot area to avoid obscuring data
    ax.legend(loc='upper center', bbox_to_anchor=(0.5, -0.08), ncol=5, fontsize=10)

    plt.tight_layout()
    plt.savefig("comparison_per_field.png", dpi=150, bbox_inches='tight', facecolor='white')
    plt.show()
    print("Saved: comparison_per_field.png")


def create_improvement_chart():
    """Create chart showing GEPA improvement over baselines."""
    fig, axes = plt.subplots(1, 2, figsize=(14, 6))

    # DSPy: Baseline vs GEPA improvement
    ax1 = axes[0]
    baseline_vals = per_field_data["DSPy Baseline"]
    gepa_vals = per_field_data["DSPy + GEPA"]
    improvements = [g - b for g, b in zip(gepa_vals, baseline_vals)]

    colors = ['#2ECC71' if imp >= 0 else '#E74C3C' for imp in improvements]
    bars1 = ax1.bar(fields, [imp * 100 for imp in improvements], color=colors, edgecolor='black', linewidth=1)
    ax1.set_ylabel('Improvement (percentage points)', fontsize=11)
    ax1.set_title('DSPy + GEPA Improvement over DSPy Baseline', fontsize=12, fontweight='bold')
    ax1.axhline(y=0, color='black', linewidth=0.8)
    ax1.set_ylim(-2, 14)
    ax1.grid(axis='y', alpha=0.3)
    plt.setp(ax1.xaxis.get_majorticklabels(), rotation=45, ha='right')

    for bar, imp in zip(bars1, improvements):
        ax1.text(bar.get_x() + bar.get_width()/2, bar.get_height() + 0.3,
                 f'+{imp:.1%}', ha='center', va='bottom', fontsize=10, fontweight='bold')

    # DSPy + BAML: Baseline vs GEPA improvement
    ax2 = axes[1]
    baml_baseline_vals = per_field_data["DSPy + BAML"]
    baml_gepa_vals = per_field_data["DSPy + BAML + GEPA"]
    baml_improvements = [g - b for g, b in zip(baml_gepa_vals, baml_baseline_vals)]

    colors2 = ['#1ABC9C' if imp >= 0 else '#E74C3C' for imp in baml_improvements]
    bars2 = ax2.bar(fields, [imp * 100 for imp in baml_improvements], color=colors2, edgecolor='black', linewidth=1)
    ax2.set_ylabel('Improvement (percentage points)', fontsize=11)
    ax2.set_title('DSPy + BAML + GEPA Improvement over DSPy + BAML', fontsize=12, fontweight='bold')
    ax2.axhline(y=0, color='black', linewidth=0.8)
    ax2.set_ylim(-2, 14)
    ax2.grid(axis='y', alpha=0.3)
    plt.setp(ax2.xaxis.get_majorticklabels(), rotation=45, ha='right')

    for bar, imp in zip(bars2, baml_improvements):
        ax2.text(bar.get_x() + bar.get_width()/2, bar.get_height() + 0.3,
                 f'+{imp:.1%}', ha='center', va='bottom', fontsize=10, fontweight='bold')

    plt.suptitle("GEPA Optimization Impact by Field", fontsize=14, fontweight='bold', y=1.02)
    plt.tight_layout()
    plt.savefig("comparison_gepa_improvement.png", dpi=150, bbox_inches='tight', facecolor='white')
    plt.show()
    print("Saved: comparison_gepa_improvement.png")


def create_dspy_vs_baml_chart():
    """Create chart comparing DSPy vs DSPy+BAML (both baseline and optimized)."""
    fig, ax = plt.subplots(figsize=(12, 6))

    x = np.arange(len(fields))
    width = 0.35

    # Compare optimized versions
    bars1 = ax.bar(x - width/2, per_field_data["DSPy + GEPA"], width,
                   label='DSPy + GEPA', color='#2ECC71', edgecolor='black', linewidth=1)
    bars2 = ax.bar(x + width/2, per_field_data["DSPy + BAML + GEPA"], width,
                   label='DSPy + BAML + GEPA', color='#1ABC9C', edgecolor='black', linewidth=1)

    ax.set_ylabel('Accuracy', fontsize=12)
    ax.set_title('Optimized Comparison: DSPy + GEPA vs DSPy + BAML + GEPA', fontsize=14, fontweight='bold')
    ax.set_xticks(x)
    ax.set_xticklabels(fields, fontsize=11)
    ax.legend(loc='lower right', fontsize=11)
    ax.set_ylim(0.7, 1.02)
    ax.grid(axis='y', alpha=0.3)

    # Add difference annotations
    for i, (v1, v2) in enumerate(zip(per_field_data["DSPy + GEPA"], per_field_data["DSPy + BAML + GEPA"])):
        diff = v2 - v1
        color = '#1ABC9C' if diff >= 0 else '#E74C3C'
        sign = '+' if diff >= 0 else ''
        ax.annotate(f'{sign}{diff:.1%}', xy=(x[i] + width/2, max(v1, v2) + 0.01),
                   ha='center', va='bottom', fontsize=9, color=color, fontweight='bold')

    plt.tight_layout()
    plt.savefig("comparison_dspy_vs_baml.png", dpi=150, bbox_inches='tight', facecolor='white')
    plt.show()
    print("Saved: comparison_dspy_vs_baml.png")


def create_summary_table():
    """Create and print a summary table."""
    print("\n" + "=" * 80)
    print("FINANCIAL ENTITY EXTRACTION - RESULTS SUMMARY")
    print("=" * 80)

    df = pd.DataFrame(results_data)
    df["Exact Match"] = df["Exact Match"].apply(lambda x: f"{x:.2%}")
    df["Mean Field"] = df["Mean Field"].apply(lambda x: f"{x:.2%}")
    print("\nOverall Metrics:")
    print(df.to_string(index=False))

    print("\n" + "-" * 80)
    print("\nPer-Field Accuracy:")
    per_field_df = pd.DataFrame(per_field_data)
    for col in per_field_df.columns[1:]:
        if per_field_df[col].iloc[0] is not None:
            per_field_df[col] = per_field_df[col].apply(lambda x: f"{x:.2%}" if x else "N/A")
    print(per_field_df.to_string(index=False))

    print("\n" + "-" * 80)
    print("\nKey Findings:")
    # Index mapping after reorder: 0=OpenAI, 1=DSPy, 2=DSPy+BAML, 3=DSPy+GEPA, 4=DSPy+BAML+GEPA
    print(f"  - Best Exact Match: DSPy + BAML + GEPA ({results_data['Exact Match'][4]:.2%})")
    print(f"  - Best Mean Field:  DSPy + GEPA ({results_data['Mean Field'][3]:.2%})")
    print(f"  - GEPA improves DSPy baseline by +{(results_data['Exact Match'][3] - results_data['Exact Match'][1])*100:.1f}pp exact match")
    print(f"  - GEPA improves DSPy+BAML baseline by +{(results_data['Exact Match'][4] - results_data['Exact Match'][2])*100:.1f}pp exact match")
    print(f"  - Biggest per-field improvement: Product (+{(per_field_data['DSPy + GEPA'][5] - per_field_data['DSPy Baseline'][5])*100:.1f}pp with GEPA)")
    print("=" * 80)


def main():
    """Generate all comparison charts."""
    print("Generating comparison charts...\n")

    # Set style
    plt.style.use('seaborn-v0_8-whitegrid')
    plt.rcParams['font.family'] = 'sans-serif'
    plt.rcParams['axes.titleweight'] = 'bold'

    create_overall_comparison_chart()
    create_per_field_comparison_chart()
    create_improvement_chart()
    create_dspy_vs_baml_chart()
    create_summary_table()

    print("\nAll charts generated successfully!")


if __name__ == "__main__":
    main()