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

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.03):
    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

# PCC_A per encoder on Track B datasets
datasets = ['deu', 'eng', 'pcm', 'swa', 'zho']

pcc_a = {
    'XLM-R-Large':  [0.4371, 0.3009, 0.5641, 0.4944, 0.2913],
    'XLM-R-Base':   [0.1400, 0.0840, 0.4719, 0.3358, 0.0982],
    'AFRO':         [0.3786, 0.3310, 0.6163, 0.4956, 0.1477],
    'mDeBERTa':     [0.1331, 0.1752, 0.4994, 0.3662, 0.1419],
    'XLM-Politics': [0.0248, 0.1749, 0.2948, 0.3827, 0.1791],
}

avg_pcc = {enc: np.mean(vals) for enc, vals in pcc_a.items()}
sorted_encoders = sorted(avg_pcc.keys(), key=lambda e: avg_pcc[e], reverse=True)

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

x = np.arange(len(datasets))
n_enc = len(sorted_encoders)
total_width = 0.75
bar_width = total_width / n_enc

enc_colors = {
    'XLM-R-Large': COLORS['valence_dark'],
    'AFRO': COLORS['valence'],
    'XLM-R-Base': COLORS['loss'],
    'mDeBERTa': COLORS['arousal'],
    'XLM-Politics': COLORS['boxes'],
}
enc_edges = {
    'XLM-R-Large': COLORS['borders'],
    'AFRO': COLORS['valence_dark'],
    'XLM-R-Base': COLORS['borders'],
    'mDeBERTa': COLORS['arousal_dark'],
    'XLM-Politics': COLORS['borders'],
}

for i, enc in enumerate(sorted_encoders):
    offset = (i - n_enc/2 + 0.5) * bar_width
    vals = [pcc_a[enc][j] for j in range(len(datasets))]
    for j, v in enumerate(vals):
        rounded_bar(ax, j + offset - bar_width*0.9/2, 0, bar_width*0.9, v,
                    enc_colors[enc], enc_edges[enc], alpha=0.6, radius=0.02)

# Danger zone line
ax.axhline(y=0.2, color=COLORS['arousal_dark'], linestyle='--', alpha=0.4, linewidth=1)

ax.set_xticks(x)
ax.set_xticklabels(['German', 'English', 'Pidgin', 'Swahili', 'Chinese'],
                    color=COLORS['borders'])
ax.set_ylabel(r'Arousal PCC $\uparrow$', color=COLORS['borders'])
ax.set_xlim(-0.5, len(datasets) - 0.5)
ax.set_ylim(0, 0.75)
ax.yaxis.grid(True, alpha=0.3, color='#CCCCCC')
ax.spines['left'].set_color(COLORS['borders'])
ax.spines['bottom'].set_color(COLORS['borders'])

# Legend
legend_elements = [
    Patch(facecolor=enc_colors[enc], edgecolor=enc_edges[enc],
          alpha=0.6, linewidth=1.5, label=f'{enc} ({avg_pcc[enc]:.2f})')
    for enc in sorted_encoders
]
legend = ax.legend(
    handles=legend_elements,
    loc='upper left',
    frameon=True,
    fancybox=True,
    framealpha=0.95,
    edgecolor='#E0E0E0',
    labelspacing=1.2,
    handletextpad=0.8,
    borderpad=0.8,
    title=r'Encoder (avg PCC$_A$)',
    title_fontsize=10,
)
legend.get_frame().set_facecolor('white')

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