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

# The toy corpus of the lesson (11 Portuguese sentences, window of 2 words).
# Left: raw co-occurrence counts. Right: PPMI, which keeps only the pairs that
# occur more often than chance. Rows and columns are grouped by subject.
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"]
order = ['gato', 'cachorro', 'bebe', 'come', 'leite', 'água', 'peixe', 'carne',
         'carro', 'automóvel', 'caminhão', 'usa', 'anda', 'gasolina', 'diesel', 'estrada',
         'corre', 'rua', 'na', 'o']
idx = {w: i for i, w in enumerate(order)}
C = np.zeros((len(order), len(order)))
for s in (t.split() for t in corpus):
    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)

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})
fig, axes = plt.subplots(1, 2, figsize=(11, 6.2))
for ax, M, title, cmap in [(axes[0], C, 'Co-occurrence counts (window 2)', 'Blues'),
                           (axes[1], PPMI, 'PPMI: counts relative to chance', 'Oranges')]:
    ax.imshow(np.ma.masked_equal(M, 0), cmap=cmap, vmin=0, vmax=M.max() * 1.05)
    for i in range(len(order)):
        for j in range(len(order)):
            if M[i, j] > 0:
                txt = f'{M[i, j]:.0f}' if M is C else f'{M[i, j]:.1f}'
                ax.text(j, i, txt, ha='center', va='center', fontsize=5.6 if M is PPMI else 6.5,
                        color='white' if M[i, j] > 0.55 * M.max() else '#30363d')
    ax.set_xticks(range(len(order)))
    ax.set_xticklabels(order, rotation=90, fontsize=7.5)
    ax.set_yticks(range(len(order)))
    ax.set_yticklabels(order, fontsize=7.5)
    for b in [7.5, 17.5]:
        ax.axhline(b, color=GRAY, lw=0.6)
        ax.axvline(b, color=GRAY, lw=0.6)
    ax.set_title(title, fontsize=9.5, loc='left', pad=18)
    ax.set_xlabel('context word', fontsize=8.5)
axes[0].set_ylabel('word', fontsize=8.5)
axes[0].text(0, 1.035, 'the largest entry of most noun rows is "o" (the)', transform=axes[0].transAxes,
             fontsize=8, color=BLUE)
axes[1].text(0, 1.035, 'now "o" is among the smallest; specific contexts dominate', transform=axes[1].transAxes,
             fontsize=8, color=ORANGE)
fig.tight_layout()

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