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

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

# Data
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]

def setup_style():
    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,
    })

# =============================================================================
# VARIANT 1: LOLLIPOP CHART
# =============================================================================
def variant_lollipop():
    setup_style()
    fig, ax = plt.subplots(figsize=(7, 4.5))
    fig.patch.set_facecolor(COLORS['bg'])
    ax.set_facecolor(COLORS['bg'])

    x = np.arange(len(languages))
    offset = 0.15

    # Stems
    for i, (sv, sa) in enumerate(zip(sigma_v, sigma_a)):
        ax.plot([i - offset, i - offset], [0, sv], color=COLORS['valence_dark'],
                linewidth=3, solid_capstyle='round')
        ax.plot([i + offset, i + offset], [0, sa], color=COLORS['arousal_dark'],
                linewidth=3, solid_capstyle='round')

    # Dots
    ax.scatter(x - offset, sigma_v, s=200, c=COLORS['valence'],
               edgecolors=COLORS['valence_dark'], linewidths=2, zorder=5,
               label=r'$\sigma_V$ (Valence)')
    ax.scatter(x + offset, sigma_a, s=200, c=COLORS['arousal'],
               edgecolors=COLORS['arousal_dark'], linewidths=2, zorder=5,
               label=r'$\sigma_A$ (Arousal)')

    # Labels
    for i, (sv, sa) in enumerate(zip(sigma_v, sigma_a)):
        ax.text(i - offset, sv + 0.06, f'{sv:.2f}', ha='center', fontsize=10,
                color=COLORS['borders'], fontweight='medium')
        ax.text(i + offset, sa + 0.06, 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_ylim(0, 1.15)
    ax.legend(loc='upper right', frameon=True, fancybox=True, framealpha=0.95)
    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('variant_lollipop.png', dpi=300, bbox_inches='tight', facecolor=COLORS['bg'])
    plt.savefig('variant_lollipop.pdf', dpi=300, bbox_inches='tight', facecolor=COLORS['bg'])
    plt.close()

# =============================================================================
# VARIANT 2: DUMBBELL CHART (shows gap directly)
# =============================================================================
def variant_dumbbell():
    setup_style()
    fig, ax = plt.subplots(figsize=(7, 4))
    fig.patch.set_facecolor(COLORS['bg'])
    ax.set_facecolor(COLORS['bg'])

    y = np.arange(len(languages))

    # Connecting lines
    for i, (sv, sa) in enumerate(zip(sigma_v, sigma_a)):
        ax.plot([sv, sa], [i, i], color=COLORS['borders'], linewidth=2, alpha=0.5, zorder=1)

    # Dots
    ax.scatter(sigma_v, y, s=250, c=COLORS['valence'], edgecolors=COLORS['valence_dark'],
               linewidths=2, zorder=5, label=r'$\sigma_V$ (Valence)')
    ax.scatter(sigma_a, y, s=250, c=COLORS['arousal'], edgecolors=COLORS['arousal_dark'],
               linewidths=2, zorder=5, label=r'$\sigma_A$ (Arousal)')

    # Labels
    for i, (sv, sa) in enumerate(zip(sigma_v, sigma_a)):
        ax.text(sv, i + 0.25, f'{sv:.2f}', ha='center', fontsize=10,
                color=COLORS['valence_dark'], fontweight='medium')
        ax.text(sa, i + 0.25, f'{sa:.2f}', ha='center', fontsize=10,
                color=COLORS['arousal_dark'], fontweight='medium')

    ax.set_yticks(y)
    ax.set_yticklabels(languages, color=COLORS['borders'])
    ax.set_xlabel(r'Learned uncertainty ($\sigma$)', color=COLORS['borders'])
    ax.set_xlim(0.2, 1.05)
    ax.legend(loc='lower right', frameon=True, fancybox=True, framealpha=0.95)
    ax.xaxis.grid(True, alpha=0.3, color='#CCCCCC')
    ax.spines['left'].set_visible(False)
    ax.spines['bottom'].set_color(COLORS['borders'])
    ax.tick_params(axis='y', length=0)

    plt.tight_layout()
    plt.savefig('variant_dumbbell.png', dpi=300, bbox_inches='tight', facecolor=COLORS['bg'])
    plt.savefig('variant_dumbbell.pdf', dpi=300, bbox_inches='tight', facecolor=COLORS['bg'])
    plt.close()

# =============================================================================
# VARIANT 3: HORIZONTAL BARS
# =============================================================================
def variant_horizontal():
    setup_style()
    fig, ax = plt.subplots(figsize=(7, 4))
    fig.patch.set_facecolor(COLORS['bg'])
    ax.set_facecolor(COLORS['bg'])

    y = np.arange(len(languages))
    height = 0.35

    bars_v = ax.barh(y - height/2, sigma_v, height, color=COLORS['valence'],
                     edgecolor=COLORS['valence_dark'], linewidth=1.5,
                     label=r'$\sigma_V$ (Valence)')
    bars_a = ax.barh(y + height/2, sigma_a, height, color=COLORS['arousal'],
                     edgecolor=COLORS['arousal_dark'], linewidth=1.5,
                     label=r'$\sigma_A$ (Arousal)')

    # Labels
    for i, (sv, sa) in enumerate(zip(sigma_v, sigma_a)):
        ax.text(sv + 0.03, i - height/2, f'{sv:.2f}', va='center', fontsize=10,
                color=COLORS['borders'], fontweight='medium')
        ax.text(sa + 0.03, i + height/2, f'{sa:.2f}', va='center', fontsize=10,
                color=COLORS['borders'], fontweight='medium')

    ax.set_yticks(y)
    ax.set_yticklabels(languages, color=COLORS['borders'])
    ax.set_xlabel(r'Learned uncertainty ($\sigma$)', color=COLORS['borders'])
    ax.set_xlim(0, 1.1)
    ax.legend(loc='lower right', frameon=True, fancybox=True, framealpha=0.95)
    ax.xaxis.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('variant_horizontal.png', dpi=300, bbox_inches='tight', facecolor=COLORS['bg'])
    plt.savefig('variant_horizontal.pdf', dpi=300, bbox_inches='tight', facecolor=COLORS['bg'])
    plt.close()

# =============================================================================
# VARIANT 4: SLOPE/BUMP CHART (shows pattern across languages)
# =============================================================================
def variant_slope():
    setup_style()
    fig, ax = plt.subplots(figsize=(7, 4.5))
    fig.patch.set_facecolor(COLORS['bg'])
    ax.set_facecolor(COLORS['bg'])

    x = np.arange(len(languages))

    # Lines
    ax.plot(x, sigma_v, color=COLORS['valence_dark'], linewidth=3,
            marker='o', markersize=12, markerfacecolor=COLORS['valence'],
            markeredgecolor=COLORS['valence_dark'], markeredgewidth=2,
            label=r'$\sigma_V$ (Valence)', zorder=5)
    ax.plot(x, sigma_a, color=COLORS['arousal_dark'], linewidth=3,
            marker='o', markersize=12, markerfacecolor=COLORS['arousal'],
            markeredgecolor=COLORS['arousal_dark'], markeredgewidth=2,
            label=r'$\sigma_A$ (Arousal)', zorder=5)

    # Fill between
    ax.fill_between(x, sigma_v, sigma_a, alpha=0.15, color=COLORS['borders'])

    # Labels
    for i, (sv, sa) in enumerate(zip(sigma_v, sigma_a)):
        offset_v = 0.07 if sv > sa else -0.07
        offset_a = 0.07 if sa > sv else -0.07
        ax.text(i, sv + offset_v, f'{sv:.2f}', ha='center', fontsize=10,
                color=COLORS['valence_dark'], fontweight='medium')
        ax.text(i, sa + offset_a, f'{sa:.2f}', ha='center', fontsize=10,
                color=COLORS['arousal_dark'], 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_ylim(0.2, 1.1)
    ax.legend(loc='upper right', frameon=True, fancybox=True, framealpha=0.95)
    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('variant_slope.png', dpi=300, bbox_inches='tight', facecolor=COLORS['bg'])
    plt.savefig('variant_slope.pdf', dpi=300, bbox_inches='tight', facecolor=COLORS['bg'])
    plt.close()

# =============================================================================
# VARIANT 5: DOT PLOT (minimal, clean)
# =============================================================================
def variant_dotplot():
    setup_style()
    fig, ax = plt.subplots(figsize=(6, 4.5))
    fig.patch.set_facecolor(COLORS['bg'])
    ax.set_facecolor(COLORS['bg'])

    y = np.arange(len(languages))

    # Dots only
    ax.scatter(sigma_v, y, s=300, c=COLORS['valence'], edgecolors=COLORS['valence_dark'],
               linewidths=2.5, zorder=5, label=r'$\sigma_V$ (Valence)')
    ax.scatter(sigma_a, y, s=300, c=COLORS['arousal'], edgecolors=COLORS['arousal_dark'],
               linewidths=2.5, zorder=5, label=r'$\sigma_A$ (Arousal)')

    # Reference lines
    for i in y:
        ax.axhline(y=i, color='#E0E0E0', linewidth=1, zorder=1)

    # Labels inside dots
    for i, (sv, sa) in enumerate(zip(sigma_v, sigma_a)):
        ax.text(sv, i, f'{sv:.2f}', ha='center', va='center', fontsize=8,
                color=COLORS['borders'], fontweight='bold')
        ax.text(sa, i, f'{sa:.2f}', ha='center', va='center', fontsize=8,
                color=COLORS['borders'], fontweight='bold')

    ax.set_yticks(y)
    ax.set_yticklabels(languages, color=COLORS['borders'], fontweight='medium')
    ax.set_xlabel(r'Learned uncertainty ($\sigma$)', color=COLORS['borders'])
    ax.set_xlim(0.25, 1.05)
    ax.set_ylim(-0.5, len(languages) - 0.5)
    ax.legend(loc='lower right', frameon=True, fancybox=True, framealpha=0.95)
    ax.spines['left'].set_visible(False)
    ax.spines['bottom'].set_color(COLORS['borders'])
    ax.tick_params(axis='y', length=0)

    plt.tight_layout()
    plt.savefig('variant_dotplot.png', dpi=300, bbox_inches='tight', facecolor=COLORS['bg'])
    plt.savefig('variant_dotplot.pdf', dpi=300, bbox_inches='tight', facecolor=COLORS['bg'])
    plt.close()

# Generate all variants
if __name__ == "__main__":
    variant_lollipop()
    print("Generated: variant_lollipop.png")

    variant_dumbbell()
    print("Generated: variant_dumbbell.png")

    variant_horizontal()
    print("Generated: variant_horizontal.png")

    variant_slope()
    print("Generated: variant_slope.png")

    variant_dotplot()
    print("Generated: variant_dotplot.png")

    print("\nAll variants saved!")
