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

def create_comprehensive_performance_chart():
    """
    Generates and saves a comprehensive performance comparison bar chart,
    including all baseline methods from the DANCE-ST paper.
    """
    # --- Data from Table 4 of the paper ---
    methods = [
        # Data-Driven
        'STAGNN', 'ATCN', 'CNN-LSTM-Attn', 'Informer', 'Graph WaveNet',
        # Physics-Informed
        'PINN-Soft',
        # Hard-Constraint
        'DC3-adapted', 'HardNet',
        # Ours
        'DANCE-ST'
    ]

    # RMSE on C-MAPSS dataset
    rmse_values = {
        'STAGNN': 17.20, 'ATCN': 18.40, 'CNN-LSTM-Attn': 16.70, 'Informer': 20.50,
        'Graph WaveNet': 16.45, 'PINN-Soft': 16.50, 'DC3-adapted': 17.95,
        'HardNet': 15.63, 'DANCE-ST': 15.63
    }
    rmse_std = {
        'STAGNN': 0.50, 'ATCN': 0.60, 'CNN-LSTM-Attn': 0.60, 'Informer': 0.70,
        'Graph WaveNet': 0.52, 'PINN-Soft': 0.55, 'DC3-adapted': 0.61,
        'HardNet': 0.51, 'DANCE-ST': 0.48
    }

    # Constraint Satisfaction (%)
    satisfaction_values = {
        'STAGNN': 85.2, 'ATCN': 87.3, 'CNN-LSTM-Attn': 89.1, 'Informer': 84.5,
        'Graph WaveNet': 90.7, 'PINN-Soft': 92.8, 'DC3-adapted': 97.5,
        'HardNet': 98.1, 'DANCE-ST': 97.2
    }

    # --- Plotting Setup ---
    plt.style.use('seaborn-v0_8-whitegrid')
    fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(20, 8))
    fig.suptitle('Comprehensive Performance Comparison on C-MAPSS Dataset', fontsize=20, weight='bold')

    # Define color groups for different categories of methods
    colors = ['#1f77b4'] * 5 + ['#ff7f0e'] * 1 + ['#2ca02c'] * 2 + ['#9467bd'] * 1
    
    # --- Plot 1: RMSE Comparison ---
    bars1 = ax1.bar(methods, [rmse_values[m] for m in methods], yerr=[rmse_std[m] for m in methods], 
                    capsize=5, color=colors)
    
    ax1.set_ylabel('RMSE (Lower is Better)', fontsize=14)
    ax1.set_xlabel('Methods', fontsize=14)
    ax1.set_title('Predictive Accuracy (RMSE)', fontsize=16, weight='bold')
    ax1.tick_params(axis='x', rotation=45, labelsize=12)
    ax1.tick_params(axis='y', labelsize=12)
    
    # Add labels to bars
    for bar, std in zip(bars1, rmse_std.values()):
        yval = bar.get_height()
        ax1.text(bar.get_x() + bar.get_width()/2.0, yval + std + 0.2, f'{yval:.2f}±{std:.2f}', 
                 ha='center', va='bottom', fontsize=11)
    ax1.set_ylim(0, max(rmse_values.values()) * 1.2)


    # --- Plot 2: Constraint Satisfaction Comparison ---
    bars2 = ax2.bar(methods, [satisfaction_values[m] for m in methods], color=colors)

    ax2.set_ylabel('Constraint Satisfaction (%) (Higher is Better)', fontsize=14)
    ax2.set_xlabel('Methods', fontsize=14)
    ax2.set_title('Constraint Satisfaction Rate', fontsize=16, weight='bold')
    ax2.tick_params(axis='x', rotation=45, labelsize=12)
    ax2.tick_params(axis='y', labelsize=12)

    # Add labels to bars
    for bar in bars2:
        yval = bar.get_height()
        ax2.text(bar.get_x() + bar.get_width()/2.0, yval + 0.5, f'{yval:.1f}%', 
                 ha='center', va='bottom', fontsize=11)
    ax2.set_ylim(80, 102) # Set y-axis to start from 80% to highlight differences

    # --- Final Touches ---
    plt.tight_layout(rect=[0, 0, 1, 0.96]) # Adjust layout to make room for suptitle
    
    # Add a legend for the color categories
    legend_patches = [
        plt.Rectangle((0,0),1,1, color='#1f77b4', label='Data-Driven'),
        plt.Rectangle((0,0),1,1, color='#ff7f0e', label='Physics-Informed (Soft)'),
        plt.Rectangle((0,0),1,1, color='#2ca02c', label='Hard-Constraint'),
        plt.Rectangle((0,0),1,1, color='#9467bd', label='Ours (Constraint-Guided)')
    ]
    fig.legend(handles=legend_patches, loc='lower center', ncol=4, fontsize=12, bbox_to_anchor=(0.5, -0.02))

    # Adjust bottom margin to make space for the legend
    plt.subplots_adjust(bottom=0.25)
    
    # Save the figure
    save_path = "performance_comparison_full.png"
    plt.savefig(save_path, dpi=300)
    print(f"Chart saved to: {save_path}")
    
if __name__ == "__main__":
    create_comprehensive_performance_chart()
