import matplotlib.pyplot as plt
import numpy as np
from io import StringIO
from matplotlib.colors import LinearSegmentedColormap
from sklearn.decomposition import NMF, TruncatedSVD
from sklearn.feature_extraction.text import TfidfVectorizer

# The lesson's 11-sentence corpus, TF-IDF, factorized into 3 components two ways:
# TruncatedSVD (the PCA family: signed components) and NMF (non-negative parts)
BLUE, RED, GRAY = '#58a6ff', '#ff7b72', 'gray'
plt.rcParams.update({'text.color': GRAY, 'axes.labelcolor': GRAY, 'axes.edgecolor': GRAY,
                     'xtick.color': GRAY, 'ytick.color': GRAY})

docs_toy = [
    "the team won the game with a late goal",
    "fans cheered when the team scored a goal",
    "the coach said the game was a fair win for the team",
    "the rocket launch put a satellite in orbit",
    "nasa delayed the rocket launch to the moon",
    "the satellite reached orbit around the moon",
    "bake the bread in a hot oven",
    "knead the dough and bake the bread slowly",
    "fresh bread from the oven smells great",
    "the team watched the rocket launch before the game",
    "astronauts ate fresh bread on the moon",
]
tv = TfidfVectorizer(stop_words='english')
T = tv.fit_transform(docs_toy)
words = tv.get_feature_names_out()
H_nmf = NMF(n_components=3, init='nndsvda', max_iter=1000, random_state=0).fit(T).components_
H_svd = TruncatedSVD(n_components=3, random_state=0).fit(T).components_

# order the words by the NMF topic they belong to, then by weight
owner = H_nmf.argmax(axis=0)
cols = np.lexsort((-H_nmf.max(axis=0), owner))
vmax = np.abs(np.vstack([H_nmf, H_svd])).max()
cmap = LinearSegmentedColormap.from_list('signed', [RED, '#e6e6e6', BLUE])

fig, axes = plt.subplots(2, 1, figsize=(10, 3.9), sharex=True)
for ax, H, title in [
        (axes[0], H_svd, f'TruncatedSVD (the PCA family): {(H_svd < -1e-9).sum()} of {H_svd.size} '
                         'weights are negative — component 2 reads "sports minus space"'),
        (axes[1], H_nmf, f'NMF: {(H_nmf < 0).sum()} negative weights — each component is one theme, '
                         'built only by adding words')]:
    im = ax.imshow(H[:, cols], cmap=cmap, vmin=-vmax, vmax=vmax, aspect='auto')
    ax.set_yticks(range(3), [f'comp. {k}' for k in range(3)], fontsize=8)
    ax.set_title(title, fontsize=9, loc='left')
axes[1].set_xticks(range(len(words)), words[cols], rotation=60, fontsize=7.5)
cb = fig.colorbar(im, ax=axes, fraction=0.025, pad=0.01)
cb.ax.tick_params(labelsize=7.5)
cb.set_label('weight of the word', fontsize=8)

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

plt.close()