import matplotlib.pyplot as plt
import numpy as np
from io import StringIO
from sklearn.datasets import fetch_20newsgroups
from sklearn.decomposition import TruncatedSVD
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.preprocessing import normalize

# Dense word vectors from co-occurrence alone (latent semantic analysis): TF-IDF of the
# 11,314 training posts of 20 Newsgroups, then a truncated SVD to 100 dimensions.
# Each word becomes a row of V * S (100 numbers). Left: cosine between word pairs,
# before (the word's sparse column of the TF-IDF matrix: which documents contain it)
# and after the SVD. Right: the nearest words of a few query words in the dense space.
news = fetch_20newsgroups(subset='train', remove=('headers', 'footers', 'quotes'))
tfidf = TfidfVectorizer(stop_words='english', min_df=5)
X = tfidf.fit_transform(news.data)
svd = TruncatedSVD(n_components=100, random_state=0).fit(X)
dense = normalize(svd.components_.T * svd.singular_values_)     # one row per word
sparse = normalize(X.T.tocsr())                                  # one row per word
idx = {w: i for i, w in enumerate(tfidf.get_feature_names_out())}

pairs = [('car', 'automobile'), ('doctor', 'physician'), ('gun', 'firearm'),
         ('orbit', 'satellite'), ('encryption', 'cryptography'),
         ('car', 'orbit'), ('doctor', 'hockey')]
before = [(sparse[idx[a]] @ sparse[idx[b]].T).toarray()[0, 0] for a, b in pairs]
after = [dense[idx[a]] @ dense[idx[b]] for a, b in pairs]

queries = ['car', 'doctor', 'orbit', 'gun', 'encryption', 'windows']
vocab = tfidf.get_feature_names_out()
neighbors = []
for q in queries:
    sim = dense @ dense[idx[q]]
    top = [j for j in np.argsort(-sim) if j != idx[q]][:5]
    neighbors.append([(vocab[j], sim[j]) for j in top])

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, (a1, a2) = plt.subplots(1, 2, figsize=(10, 4.2), gridspec_kw={'width_ratios': [1, 1]})

y = np.arange(len(pairs))[::-1]
a1.barh(y + 0.2, before, height=0.38, color='gray', label='sparse: which documents contain the word')
a1.barh(y - 0.2, after, height=0.38, color=BLUE, label='dense: 100-d vector after the SVD')
for yi, b, a in zip(y, before, after):
    a1.text(max(b, 0) + 0.015, yi + 0.2, f'{b:.2f}', va='center', fontsize=7.8)
    a1.text(max(a, 0) + 0.015, yi - 0.2, f'{abs(a) if abs(a) < 0.005 else a:.2f}', va='center', fontsize=7.8, color=BLUE)
a1.set_yticks(y)
a1.set_yticklabels([f'{a} / {b}' for a, b in pairs], fontsize=8.5)
a1.axhline(1.5, color='gray', lw=0.6, ls=':')
a1.set_xlim(0, 1.12)
a1.set_xlabel('cosine similarity between the two words')
a1.set_title('Related words move together;\nunrelated ones stay at zero', fontsize=9.5)
a1.legend(fontsize=7.5, frameon=False, loc='lower right')
for s in ['top', 'right']:
    a1.spines[s].set_visible(False)

a2.axis('off')
a2.set_xlim(0, 1)
a2.set_ylim(0, 1)
a2.set_title('Nearest words in the dense space (cosine)', fontsize=9.5)
colors = [BLUE, ORANGE, GREEN, RED, BLUE, ORANGE]
for r, (q, nb) in enumerate(zip(queries, neighbors)):
    yy = 0.92 - r * 0.165
    a2.text(0.0, yy, q, fontsize=9, color=colors[r], fontweight='bold', va='center')
    line = ',  '.join(f'{w} {c:.2f}' for w, c in nb[:3])
    line2 = ',  '.join(f'{w} {c:.2f}' for w, c in nb[3:])
    a2.text(0.24, yy + 0.035, line, fontsize=8, va='center')
    a2.text(0.24, yy - 0.035, line2, fontsize=8, va='center')

plt.tight_layout()

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

plt.close()