import matplotlib.pyplot as plt
import numpy as np
from matplotlib.patches import FancyBboxPatch

COLORS = {
    'valence': '#C2E3F1',
    'arousal': '#FFA5A3',
    'boxes': '#F7F1E4',
    'loss': '#FFE09A',
    'borders': '#4F6A72',
    'valence_dark': '#8BBDD9',
    'arousal_dark': '#E87D7A',
    'bg': '#FAFAFA',
}

plt.rcParams.update({
    'font.family': 'sans-serif',
    'font.sans-serif': ['Helvetica Neue', 'Arial', 'DejaVu Sans'],
    'font.size': 11,
    'axes.labelsize': 13,
    'xtick.labelsize': 11,
    'ytick.labelsize': 10,
    'legend.fontsize': 10,
    'axes.spines.top': False,
    'axes.spines.right': False,
})

def rounded_bar(ax, x, y, width, height, color, edgecolor, alpha=0.6, radius=0.06):
    box = FancyBboxPatch((x, y), width, height,
                         boxstyle="square,pad=0",
                         facecolor=color, edgecolor=edgecolor,
                         linewidth=1.5, alpha=alpha, zorder=3)
    ax.add_patch(box)
    return box

# Data from MEETING_EXPERIMENTS.md — Pidgin progression
labels = [
    'XLM-R-Large\n(global)',
    'Tw-RoBERTa\n(PCM only)',
    'Tw-RoBERTa\n+ Uncertainty\n(PCM separate)',
]
rmse = [0.9848, 0.7740, 0.6238]

# Gradient: arousal (worst) -> valence -> valence_dark (best)
bar_colors = [COLORS['arousal'], COLORS['valence'], COLORS['valence_dark']]
edge_colors = [COLORS['arousal_dark'], COLORS['valence_dark'], COLORS['borders']]

fig, ax = plt.subplots(figsize=(7, 4.5))
fig.patch.set_facecolor(COLORS['bg'])
ax.set_facecolor(COLORS['bg'])

x = np.arange(len(labels))
bar_width = 0.5

for i, (val, fc, ec) in enumerate(zip(rmse, bar_colors, edge_colors)):
    rounded_bar(ax, i - bar_width/2, 0, bar_width, val, fc, ec, alpha=0.6)

# Value labels
for i, val in enumerate(zip(x, rmse)):
    ax.text(val[0], val[1] + 0.03, f'{val[1]:.2f}',
            ha='center', fontsize=10, color=COLORS['borders'], fontweight='medium')

ax.set_xticks(x)
ax.set_xticklabels(labels, color=COLORS['borders'])
ax.set_ylabel(r'RMSE $\downarrow$', color=COLORS['borders'])
ax.set_xlim(-0.5, len(labels) - 0.5)
ax.set_ylim(0, 1.15)
ax.yaxis.grid(True, alpha=0.3, color='#CCCCCC')
ax.spines['left'].set_color(COLORS['borders'])
ax.spines['bottom'].set_color(COLORS['borders'])

plt.tight_layout()
plt.savefig('figures/pidgin_progression.pdf', dpi=300, bbox_inches='tight', facecolor=COLORS['bg'])
plt.savefig('figures/pidgin_progression.png', dpi=300, bbox_inches='tight', facecolor=COLORS['bg'])
print("Saved: pidgin_progression.pdf and .png")
