#First, Check out the average confidence of each method for each dataset
import matplotlib.pyplot as plt
import numpy as np

# --- Data Setup ---
# Categories for the x-axis
categories = [
    'AIME 2024\n(Pass@1)', 'Codeforces\n(Percentile)', 'GPQA Diamond\n(Pass@1)',
    'MATH-500\n(Pass@1)', 'MMLU\n(Pass@1)', 'SWE-bench Verified\n(Resolved)'
]

# Labels for each bar group (from the legend)
model_labels = [
    'DeepSeek-R1', 'OpenAI-o1-1217', 'DeepSeek-R1-32B',
    'OpenAI-o1-mini', 'DeepSeek-V3'
]

# Data for each model across the categories.
# I've used the data from your image. I used None for missing bars.
data = {
    'DeepSeek-R1': [79.8, 96.3, 71.5, 97.3, 90.8, 49.2],
    'OpenAI-o1-1217': [79.2, 96.6, 75.7, 96.4, 91.8, 48.9],
    'DeepSeek-R1-32B': [72.6, 90.6, 62.1, 90.0, 87.4, 36.8],
    'OpenAI-o1-mini': [63.6, 93.4, 60.0, 90.2, 85.2, 41.6],
    'DeepSeek-V3': [39.2, 58.7, 59.1, 94.3, 88.5, 42.0],
}


# --- Chart Styling ---
# Define colors and hatches to match the reference image
colors = ['#FF7F00', '#0000FF', '#E41A1C', '#A65628', '#4DAF4A']
hatches = ['', '////', '', '', '']


# --- Plotting ---
# Create the figure and axes objects
def create_plot(data, data_labels, categories,
                colors = ['#FF7F00', '#0000FF', '#E41A1C', '#A65628', '#4DAF4A'],
                hatches = ['', '', '', '', ''],hlines = None):
    
    fig, ax = plt.subplots(figsize=(12, 7))

    if hlines is not None:
        hline_colors = ['#005f73','#0a9396','#43aa8b','#495057','#ffc300','#ffd60a']
        hline_captions = [0.9, 0.7, 0.5, 0.3, 0.1, 0]
        for i,hline in enumerate(hlines):
            ax.axhline(hline, color=hline_colors[i], linestyle='--', linewidth=1,label = f"Threshold: {hline_captions[i]}")

    # Set the positions and width for the bars
    n_data = len(data_labels)
    bar_width = 0.15
    x = np.arange(len(categories))  # the label locations

    # Loop through each model to plot its bars
    for i, data_label in enumerate(data_labels):
        # Calculate the position of each bar
        position = x - (bar_width * n_data / 2) + (i * bar_width)

        # Plot the bar
        bars = ax.bar(position, data[data_label], bar_width,
                        label=data_label,
                        color=colors[i],
                        hatch=hatches[i],
                        edgecolor='grey',
                        linewidth=0.5)

        # Add data labels on top of each bar
        for bar in bars:
            height = bar.get_height()
            if height is not None:
                ax.annotate(f'{height}',
                        xy=(bar.get_x() + bar.get_width() / 2, height),
                        xytext=(0, 3),  # 3 points vertical offset
                        textcoords="offset points",
                        ha='center', va='bottom',
                        fontsize=9)


    # --- Customization ---
    # Set the y-axis label
    ax.set_ylabel('Average Explanation Confidence (%)', fontsize=12)
    ax.set_ylim(0, 1) # Give some space for the top labels

    # Set the x-axis ticks and labels
    ax.set_xticks(x)
    ax.set_xticklabels(categories, fontsize=10)

    # Add a faint horizontal grid
    ax.yaxis.grid(True, linestyle='--', which='major', color='grey', alpha=0.3)
    ax.set_axisbelow(True) # Send grid lines to the back

    # Remove top and right spines
    ax.spines['top'].set_visible(False)
    ax.spines['right'].set_visible(False)
    ax.spines['left'].set_color('grey')
    ax.spines['bottom'].set_color('grey')


    # --- Legend ---
    # Create a legend and place it horizontally on top of the chart
    ax.legend(
        loc='upper center',
        bbox_to_anchor=(0.5, 1.15), # Position above the plot
        ncol=len(data_labels),   # Display in a single row
        frameon=False,            # Remove the frame
        fontsize=12
    )

    # Adjust layout to prevent labels from being cut off
    plt.tight_layout()

    # Display the plot
    plt.show()
