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

# How a sentence-transformer turns a sentence into one vector (a "bi-encoder").
# Each sentence goes separately through the same network: tokens -> contextual
# token vectors -> mean pooling -> unit length. Two sentences are compared by
# the cosine of their vectors. Bottom: how the network was trained.
# The subword tokens shown are the real ones of all-MiniLM-L6-v2 ("lender" -> lend, ##er).
GRAY = '#8b949e'
BLUE, ORANGE, GREEN, RED = '#58a6ff', '#f0883e', '#3fb950', '#ff7b72'
plt.rcParams.update({'text.color': GRAY})
fig, ax = plt.subplots(figsize=(10, 5.3))
ax.set_xlim(0, 10)
ax.set_ylim(0, 5.3)
ax.axis('off')


def box(x, y, w, h, text, color, fill=0.15, fs=8.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.1))
    ax.text(x, y, text, ha='center', va='center', fontsize=fs, color=color)


def arrow(x0, y0, x1, y1, color=GRAY):
    ax.annotate('', xy=(x1, y1), xytext=(x0, y0), arrowprops=dict(arrowstyle='->', color=color, lw=1))


for row, (sent, toks, col, name) in enumerate([
        ('"she went to the bank"', ['she', 'went', 'to', 'the', 'bank'], BLUE, 'u'),
        ('"he visited a lender"', ['he', 'visited', 'a', 'lend', '##er'], ORANGE, 'v')]):
    y = 4.35 - 1.55 * row
    ax.text(0.1, y, sent, fontsize=9, va='center', color=col)
    arrow(1.95, y, 2.25, y)
    for k, t in enumerate(toks):
        box(2.55 + 0.42 * k, y, 0.38, 0.34, t, col, fs=6.8)
    arrow(4.65, y, 4.95, y)
    box(5.75, y, 1.5, 0.62, 'transformer\n(same weights)', GRAY, fill=0.1, fs=8)
    arrow(6.55, y, 6.85, y)
    for k in range(5):
        ax.add_patch(plt.Rectangle((6.95 + 0.13 * k, y - 0.28), 0.09, 0.56, fc=col, alpha=0.35 + 0.1 * k, ec='none'))
    arrow(7.65, y, 7.95, y)
    box(8.45, y, 0.85, 0.5, 'mean\n+ unit length', col, fs=7.2)
    ax.add_patch(plt.Rectangle((9.05, y - 0.28), 0.12, 0.56, fc=col, alpha=0.8, ec='none'))
    ax.text(9.3, y, name, fontsize=11, va='center', color=col, fontweight='bold')

ax.text(2.55 + 0.84, 5.0, 'subword tokens', ha='center', fontsize=8)
ax.text(7.2, 5.0, 'one vector per token,\ncomputed from the whole sentence', ha='center', fontsize=7.5)
ax.text(9.15, 5.0, '384 numbers', ha='center', fontsize=7.5)
ax.annotate('', xy=(9.6, 2.9), xytext=(9.6, 4.25), arrowprops=dict(arrowstyle='<->', color=GREEN, lw=1.2))
ax.text(9.72, 3.58, 'cos(u, v)\n= u · v', fontsize=8.5, color=GREEN, va='center')

# training
ax.plot([0.1, 9.9], [1.95, 1.95], color=GRAY, lw=0.5, alpha=0.5)
ax.text(0.1, 1.62, 'How the network was trained (Sentence-BERT)', fontsize=9.5, color=GRAY)
ax.text(0.1, 1.1, 'pairs that mean the same thing', fontsize=8.5, color=GREEN)
ax.text(0.1, 0.8, '("a man plays guitar", "someone strums an instrument")', fontsize=7.8, color=GRAY)
ax.text(4.3, 0.95, '→  push cos(u, v) up', fontsize=8.5, color=GREEN)
ax.text(0.1, 0.42, 'pairs that do not', fontsize=8.5, color=RED)
ax.text(0.1, 0.12, '("a man plays guitar", "the market fell today")', fontsize=7.8, color=GRAY)
ax.text(4.3, 0.27, '→  push cos(u, v) down', fontsize=8.5, color=RED)
ax.text(6.6, 0.95, 'Starts from a transformer pretrained on\nplain text; this second training makes\n'
                   'the cosine track meaning, sentence by\nsentence.', fontsize=7.8, color=GRAY, va='center')

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