import matplotlib.pyplot as plt
import numpy as np
from io import StringIO
from sklearn.decomposition import TruncatedSVD

# The toy corpus of the lesson, window of 2 words. Each word's row of the
# word x context matrix is compressed to 2 numbers by a truncated SVD.
# Left: SVD of the raw counts. Right: SVD of the PPMI matrix.
corpus = ["o gato bebe leite", "o gato come peixe",
          "o cachorro bebe água", "o cachorro come carne",
          "o cachorro corre na rua",
          "o carro usa gasolina", "o carro corre na estrada",
          "o automóvel usa gasolina", "o automóvel anda na estrada",
          "o caminhão usa diesel", "o caminhão corre na estrada"]
sentences = [s.split() for s in corpus]
vocab = sorted({w for s in sentences for w in s})
idx = {w: i for i, w in enumerate(vocab)}
C = np.zeros((len(vocab), len(vocab)))
for s in sentences:
    for i, w in enumerate(s):
        for j in range(max(0, i - 2), min(len(s), i + 3)):
            if j != i:
                C[idx[w], idx[s[j]]] += 1
N = C.sum()
with np.errstate(divide='ignore'):
    PPMI = np.maximum(np.log2(C * N / np.outer(C.sum(1), C.sum(0))), 0)

animals = {'gato', 'cachorro', 'bebe', 'come', 'leite', 'água', 'peixe', 'carne'}
vehicles = {'carro', 'automóvel', 'caminhão', 'usa', 'anda', 'gasolina', 'diesel', 'estrada'}
GRAY = '#8b949e'
BLUE, ORANGE, GREEN, RED = '#58a6ff', '#f0883e', '#3fb950', '#ff7b72'
plt.rcParams.update({'text.color': GRAY, 'axes.labelcolor': GRAY, 'axes.edgecolor': GRAY,
                     'xtick.color': GRAY, 'ytick.color': GRAY})


def cos(M, a, b):
    u, v = M[idx[a]], M[idx[b]]
    return u @ v / (np.linalg.norm(u) * np.linalg.norm(v))


labels = {}
fig, axes = plt.subplots(1, 2, figsize=(10.5, 4.9))
for ax, M, title in [(axes[0], C, 'SVD of the raw counts'), (axes[1], PPMI, 'SVD of the PPMI matrix')]:
    Z = TruncatedSVD(n_components=2, random_state=0).fit_transform(M)
    for w in vocab:
        col = GREEN if w in animals else BLUE if w in vehicles else GRAY
        x, y = Z[idx[w]]
        ax.scatter(x, y, s=22, color=col, zorder=3)
    labels[ax] = [(w, *Z[idx[w]], GREEN if w in animals else BLUE if w in vehicles else GRAY)
                  for w in vocab]
    ax.axhline(0, color=GRAY, lw=0.5, alpha=0.5)
    ax.axvline(0, color=GRAY, lw=0.5, alpha=0.5)
    ax.set_xlabel('dimension 1', fontsize=8.5)
    ax.set_ylabel('dimension 2', fontsize=8.5)
    ax.tick_params(labelsize=7.5)
    ax.set_title(title, fontsize=9.5, loc='left')
    ax.text(0.02, 0.03, f"cos(gato, carro) = {cos(Z, 'gato', 'carro'):.2f}\n"
                        f"cos(gato, cachorro) = {cos(Z, 'gato', 'cachorro'):.2f}",
            transform=ax.transAxes, fontsize=8, color=ORANGE, va='bottom')
    x0, x1 = ax.get_xlim()
    ax.set_xlim(x0, x1 + 0.12 * (x1 - x0))
fig.tight_layout()

# place labels in display space: start right of each point, push overlapping
# labels apart vertically, then draw a thin leader line when a label moved
for ax, items in labels.items():
    tr = ax.transData
    pts = np.array([tr.transform((x, y)) for _, x, y, _ in items])
    pos = pts + [6, 0]
    w_px = np.array([len(w) * 6.2 for w, *_ in items])
    for _ in range(400):
        moved = False
        for i in range(len(pos)):
            for j in range(i + 1, len(pos)):
                if abs(pos[i, 0] - pos[j, 0]) < (w_px[i] + w_px[j]) / 2 + 2 and abs(pos[i, 1] - pos[j, 1]) < 12:
                    d = 1.0 if pos[i, 1] >= pos[j, 1] else -1.0
                    pos[i, 1] += d * 1.5
                    pos[j, 1] -= d * 1.5
                    moved = True
        if not moved:
            break
    inv = tr.inverted()
    for (w, x, y, col), p, q in zip(items, pos, pts):
        tx, ty = inv.transform(p)
        ax.annotate(w, (x, y), xytext=(tx, ty), fontsize=7.8, color=col, va='center', ha='left',
                    arrowprops=dict(arrowstyle='-', color=col, lw=0.5, alpha=0.6)
                    if np.hypot(*(p - q - [6, 0])) > 4 else None)

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