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

# Semantic search is nearest-neighbor search over embeddings.
# Offline (once): encode every document into a unit vector -> an N x d matrix.
# Online (per query): encode the query, one matrix-vector product gives all the
# cosines, keep the k largest. Right: the same thing drawn in 2-D, where the
# cosine is the angle between the query and each document.
GRAY = '#8b949e'
BLUE, ORANGE, GREEN, RED = '#58a6ff', '#f0883e', '#3fb950', '#ff7b72'
plt.rcParams.update({'text.color': GRAY, 'axes.edgecolor': GRAY})
fig = plt.figure(figsize=(10, 4.4))
ax = fig.add_axes([0, 0, 0.62, 1])
ax.set_xlim(0, 6.2)
ax.set_ylim(0, 4.4)
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))


ax.text(0.1, 4.15, 'Offline, once', fontsize=9.5, color=BLUE)
box(0.75, 3.4, 1.2, 0.9, 'N documents\n(FAQ, posts,\nproducts)', BLUE, fs=7.8)
arrow(1.4, 3.4, 1.75, 3.4)
box(2.3, 3.4, 0.95, 0.5, 'encoder', GRAY, fill=0.1)
arrow(2.8, 3.4, 3.15, 3.4)
for k in range(6):
    ax.add_patch(plt.Rectangle((3.25, 3.78 - 0.14 * k), 1.1, 0.1, fc=BLUE, alpha=0.3 + 0.1 * k, ec='none'))
ax.text(3.8, 2.72, 'N × 384 matrix\n(the "index")', ha='center', fontsize=7.8, color=BLUE)

ax.text(0.1, 2.25, 'Online, per query', fontsize=9.5, color=ORANGE)
box(0.75, 1.5, 1.2, 0.5, '"esqueci a senha"', ORANGE, fs=7.8)
arrow(1.4, 1.5, 1.75, 1.5)
box(2.3, 1.5, 0.95, 0.5, 'encoder', GRAY, fill=0.1)
arrow(2.8, 1.5, 3.15, 1.5)
ax.add_patch(plt.Rectangle((3.25, 1.45), 1.1, 0.1, fc=ORANGE, alpha=0.9, ec='none'))
ax.text(3.8, 1.2, 'query vector q', ha='center', fontsize=7.8, color=ORANGE)
arrow(4.45, 1.5, 4.8, 2.3)
arrow(4.45, 3.4, 4.8, 2.6)
box(5.35, 2.45, 0.95, 0.75, 'index · q\n= N cosines\n→ top k', GREEN, fs=7.8)
ax.text(5.35, 1.65, 'a matrix product;\nfor millions of rows,\na vector database', ha='center',
        fontsize=7.3, color=GRAY, va='top')

# right: the geometry
ax2 = fig.add_axes([0.66, 0.1, 0.31, 0.8])
rng = np.random.default_rng(3)
ang = np.concatenate([rng.normal(0.5, 0.12, 5), rng.normal(1.9, 0.15, 5), rng.normal(3.4, 0.15, 5)])
cols = [BLUE] * 15
qa = 0.62
for a in ang:
    ax2.plot([0, np.cos(a)], [0, np.sin(a)], color=GRAY, lw=0.5, alpha=0.4)
    ax2.plot(np.cos(a), np.sin(a), 'o', color=BLUE, ms=5)
ax2.annotate('', xy=(np.cos(qa), np.sin(qa)), xytext=(0, 0), arrowprops=dict(arrowstyle='->', color=ORANGE, lw=1.8))
ax2.text(np.cos(qa) + 0.05, np.sin(qa) + 0.08, 'query', color=ORANGE, fontsize=8.5)
near = np.argsort(np.abs(ang - qa))[:3]
for j in near:
    ax2.plot(np.cos(ang[j]), np.sin(ang[j]), 'o', color=GREEN, ms=8, mfc='none', mew=1.5)
t = np.linspace(0, 2 * np.pi, 200)
ax2.plot(np.cos(t), np.sin(t), color=GRAY, lw=0.5, alpha=0.4)
ax2.text(0, -1.3, 'unit vectors on a circle: the k nearest\nare the smallest angles (green)', ha='center',
         fontsize=7.8)
ax2.set_xlim(-1.25, 1.25)
ax2.set_ylim(-1.45, 1.25)
ax2.set_aspect('equal')
ax2.axis('off')

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