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):
    """Draw a single rounded-corner bar with glassy look."""
    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 — Learned sigma per language (seed 42)
languages = ['English', 'German', 'Pidgin', 'Chinese']
sigma_v = [0.91, 0.52, 0.91, 0.39]
sigma_a = [0.42, 0.79, 0.42, 0.47]

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

x = np.arange(len(languages))
bar_width = 0.32
gap = 0.04

# Draw rounded bars
for i, (sv, sa) in enumerate(zip(sigma_v, sigma_a)):
    # Valence bar (left)
    rounded_bar(ax, i - bar_width - gap/2, 0, bar_width, sv,
                COLORS['valence'], COLORS['valence_dark'], alpha=0.6)
    # Arousal bar (right)
    rounded_bar(ax, i + gap/2, 0, bar_width, sa,
                COLORS['arousal'], COLORS['arousal_dark'], alpha=0.6)

# Value labels
for i, (sv, sa) in enumerate(zip(sigma_v, sigma_a)):
    ax.text(i - bar_width/2 - gap/2, sv + 0.04, f'{sv:.2f}',
            ha='center', fontsize=10, color=COLORS['borders'], fontweight='medium')
    ax.text(i + bar_width/2 + gap/2, sa + 0.04, f'{sa:.2f}',
            ha='center', fontsize=10, color=COLORS['borders'], fontweight='medium')

ax.set_xticks(x)
ax.set_xticklabels(languages, color=COLORS['borders'])
ax.set_ylabel(r'Learned uncertainty ($\sigma$)', color=COLORS['borders'])
ax.set_xlim(-0.6, len(languages) - 0.4)
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'])

# Legend with colored patches
from matplotlib.patches import Patch
legend_elements = [
    Patch(facecolor=COLORS['valence'], edgecolor=COLORS['valence_dark'],
          alpha=0.6, linewidth=1.5, label=r'$\sigma_V$ (Valence)'),
    Patch(facecolor=COLORS['arousal'], edgecolor=COLORS['arousal_dark'],
          alpha=0.6, linewidth=1.5, label=r'$\sigma_A$ (Arousal)'),
]
legend = ax.legend(
    handles=legend_elements,
    loc='upper right',
    frameon=True,
    fancybox=True,
    framealpha=0.95,
    edgecolor='#E0E0E0',
    labelspacing=1.2,
    handletextpad=0.8,
    borderpad=0.8,
)
legend.get_frame().set_facecolor('white')

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