import matplotlib.pyplot as plt
from io import StringIO
from matplotlib.patches import FancyBboxPatch

# How word2vec turns running text into training examples. Top: a window of
# 2 words around the center word "bebe". Bottom: the examples that window
# produces for skip-gram, for CBOW, and with negative sampling.
GRAY = '#8b949e'
BLUE, ORANGE, GREEN, RED = '#58a6ff', '#f0883e', '#3fb950', '#ff7b72'
plt.rcParams.update({'text.color': GRAY})
fig, ax = plt.subplots(figsize=(10, 4.6))
ax.set_xlim(0, 10)
ax.set_ylim(-0.1, 5.2)
ax.axis('off')


def box(x, y, text, color, w=1.05, h=0.42, fill=0.18, fs=9.5):
    ax.add_patch(FancyBboxPatch((x - w / 2, y - h / 2), w, h, boxstyle='round,pad=0.02,rounding_size=0.08',
                                fc=color, alpha=fill, ec=color, lw=1.2))
    ax.text(x, y, text, ha='center', va='center', fontsize=fs, color=color)


words = ['o', 'gato', 'bebe', 'leite', 'gelado', 'hoje']
xs = [1.4 + 1.3 * i for i in range(len(words))]
for i, (w, x) in enumerate(zip(words, xs)):
    col = ORANGE if i == 2 else BLUE if abs(i - 2) <= 2 else GRAY
    box(x, 4.55, w, col)
ax.annotate('', xy=(xs[0] - 0.55, 4.05), xytext=(xs[4] + 0.55, 4.05),
            arrowprops=dict(arrowstyle='|-|', color=BLUE, lw=1, mutation_scale=4))
ax.text((xs[0] + xs[4]) / 2, 3.8, 'window = 2 words on each side of the center word', ha='center',
        fontsize=8.5, color=BLUE)
ax.text(xs[2], 5.02, 'center word', ha='center', fontsize=8.5, color=ORANGE)
ax.text(xs[5], 5.02, 'outside', ha='center', fontsize=8.5, color=GRAY)

# skip-gram
ax.text(0.3, 3.45, 'Skip-gram\nthe center word predicts each neighbor', fontsize=9, color=GRAY, va='top')
box(0.95, 1.8, 'bebe', ORANGE)
for k, w in enumerate(['o', 'gato', 'leite', 'gelado']):
    y = 2.55 - 0.5 * k
    box(2.75, y, w, BLUE)
    ax.annotate('', xy=(2.2, y), xytext=(1.5, 1.8), arrowprops=dict(arrowstyle='->', color=GRAY, lw=0.9))
ax.text(1.8, 0.3, '4 training examples:\n(bebe → o), (bebe → gato), ...', fontsize=8, color=GRAY, ha='center')

# CBOW
ax.text(3.75, 3.45, 'CBOW\nthe neighbors predict the center word', fontsize=9, color=GRAY, va='top')
for k, w in enumerate(['o', 'gato', 'leite', 'gelado']):
    y = 2.55 - 0.5 * k
    box(4.35, y, w, BLUE)
    ax.annotate('', xy=(5.75, 1.8), xytext=(4.9, y), arrowprops=dict(arrowstyle='->', color=GRAY, lw=0.9))
box(6.3, 1.8, 'bebe', ORANGE)
ax.text(5.3, 0.3, '1 training example:\n(average of the 4 → bebe)', fontsize=8, color=GRAY, ha='center')

# negative sampling
ax.text(7.15, 3.45, 'Negative sampling\nreal pairs against random pairs', fontsize=9, color=GRAY, va='top')
rows = [('bebe', 'leite', 'real pair', '→ 1', GREEN),
        ('bebe', 'diesel', 'random word', '→ 0', RED),
        ('bebe', 'estrada', 'random word', '→ 0', RED)]
for k, (a, b, lab, t, col) in enumerate(rows):
    y = 2.45 - 0.55 * k
    box(7.55, y, a, ORANGE, w=0.75, fs=8.5)
    box(8.4, y, b, col, w=0.85, fs=8.5)
    ax.text(8.95, y, t, va='center', fontsize=9.5, color=col)
    ax.text(9.35, y, lab, va='center', fontsize=7.5, color=GRAY)
ax.text(8.3, 0.2, 'a yes/no question per pair:\n"did these two words really\nappear together?"', fontsize=8,
        color=GRAY, ha='center')

buffer = StringIO()
plt.savefig(buffer, format="svg", transparent=True)
print(buffer.getvalue())
plt.close()