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.04):
    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 — English vs Pidgin sigma convergence, all 3 seeds
seeds = ['Seed 42', 'Seed 12', 'Seed 73']

eng_sigma_v = [0.9116, 0.6206, 0.3056]
eng_sigma_a = [0.4177, 0.5823, 0.4568]
pcm_sigma_v = [0.9105, 0.6196, 0.3055]
pcm_sigma_a = [0.4175, 0.5811, 0.4566]

fig, axes = plt.subplots(1, 3, figsize=(10, 4.5), sharey=True)
fig.patch.set_facecolor(COLORS['bg'])

bar_width = 0.28
gap = 0.04

for i, (seed, ax) in enumerate(zip(seeds, axes)):
    ax.set_facecolor(COLORS['bg'])

    eng_vals = [eng_sigma_v[i], eng_sigma_a[i]]
    pcm_vals = [pcm_sigma_v[i], pcm_sigma_a[i]]
    dim_colors = [COLORS['valence'], COLORS['arousal']]
    dim_edges = [COLORS['valence_dark'], COLORS['arousal_dark']]

    for j in range(2):
        # English bar (solid, left)
        rounded_bar(ax, j - bar_width - gap/2, 0, bar_width, eng_vals[j],
                    dim_colors[j], dim_edges[j], alpha=0.6)
        # Pidgin bar (slightly more transparent, right)
        rounded_bar(ax, j + gap/2, 0, bar_width, pcm_vals[j],
                    dim_colors[j], dim_edges[j], alpha=0.35)

    # Delta labels
    for j, (ev, pv) in enumerate(zip(eng_vals, pcm_vals)):
        delta = abs(ev - pv)
        y_pos = max(ev, pv) + 0.04
        ax.text(j, y_pos, f'{delta:.4f}', ha='center', fontsize=8,
                color=COLORS['borders'], fontweight='medium')

    ax.set_xticks([0, 1])
    ax.set_xticklabels([r'$\sigma^2_V$', r'$\sigma^2_A$'], color=COLORS['borders'])
    ax.set_title(seed, fontsize=12, color=COLORS['borders'], fontweight='medium')
    ax.set_xlim(-0.5, 1.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'])

axes[0].set_ylabel(r'Learned variance ($\sigma^2$)', color=COLORS['borders'])

# Legend — solid vs transparent to distinguish English/Pidgin
legend_elements = [
    Patch(facecolor='#DDDDDD', edgecolor=COLORS['borders'],
          alpha=0.6, linewidth=1.2, label='English'),
    Patch(facecolor='#DDDDDD', edgecolor=COLORS['borders'],
          alpha=0.35, linewidth=1.2, label='Pidgin'),
]
fig.legend(handles=legend_elements, loc='upper center', ncol=2,
           frameon=True, fancybox=True, framealpha=0.95,
           edgecolor='#E0E0E0', bbox_to_anchor=(0.5, 1.02), fontsize=11)

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