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.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 — Translation vs native model, Russian restaurant
metrics = [r'PCC$_V$ $\uparrow$', r'PCC$_A$ $\uparrow$', r'RMSE $\downarrow$']
translated = [0.938, 0.557, 0.9651]
native = [0.934, 0.586, 0.9158]

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

x = np.arange(len(metrics))
bar_width = 0.3
gap = 0.04

for i, (tv, nv) in enumerate(zip(translated, native)):
    # Translated bar (loss yellow, left)
    rounded_bar(ax, i - bar_width - gap/2, 0, bar_width, tv,
                COLORS['loss'], COLORS['borders'], alpha=0.6)
    # Native bar (valence blue, right)
    rounded_bar(ax, i + gap/2, 0, bar_width, nv,
                COLORS['valence'], COLORS['valence_dark'], alpha=0.6)

# Value labels
for i, (tv, nv) in enumerate(zip(translated, native)):
    ax.text(i - bar_width/2 - gap/2, tv + 0.03, f'{tv:.2f}',
            ha='center', fontsize=10, color=COLORS['borders'], fontweight='medium')
    ax.text(i + bar_width/2 + gap/2, nv + 0.03, f'{nv:.2f}',
            ha='center', fontsize=10, color=COLORS['borders'], fontweight='medium')

ax.set_xticks(x)
ax.set_xticklabels(metrics, color=COLORS['borders'])
ax.set_xlim(-0.55, len(metrics) - 0.45)
ax.set_ylim(0, 1.2)
ax.yaxis.grid(True, alpha=0.3, color='#CCCCCC')
ax.spines['left'].set_color(COLORS['borders'])
ax.spines['bottom'].set_color(COLORS['borders'])

legend_elements = [
    Patch(facecolor=COLORS['loss'], edgecolor=COLORS['borders'],
          alpha=0.6, linewidth=1.5, label=r'Translated (RU$\rightarrow$EN)'),
    Patch(facecolor=COLORS['valence'], edgecolor=COLORS['valence_dark'],
          alpha=0.6, linewidth=1.5, label='Native (ruRoBERTa)'),
]
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/translation_vs_native.pdf', dpi=300, bbox_inches='tight', facecolor=COLORS['bg'])
plt.savefig('figures/translation_vs_native.png', dpi=300, bbox_inches='tight', facecolor=COLORS['bg'])
print("Saved: translation_vs_native.pdf and .png")
