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

# The three-sentence corpus of the lesson: the same matrix twice, as raw counts
# and as TF-IDF weights (scikit-learn defaults: smoothed idf, L2-normalized rows).
docs = ["O gato comeu o peixe.",
        "O cachorro comeu a ração do gato.",
        "Gato dorme, gato sonha."]

cv = CountVectorizer()
C = cv.fit_transform(docs).toarray()
tv = TfidfVectorizer()
T = tv.fit_transform(docs).toarray()
vocab = tv.get_feature_names_out()
df = (C > 0).sum(axis=0)

BLUE, ORANGE = '#58a6ff', '#f0883e'
plt.rcParams.update({'text.color': 'gray', 'axes.labelcolor': 'gray', 'axes.edgecolor': 'gray',
                     'xtick.color': 'gray', 'ytick.color': 'gray'})
# opaque cells (light gray -> color) so the black numbers stay readable on a dark page too
cmaps = [LinearSegmentedColormap.from_list('c', ['#eef1f4', ORANGE]),
         LinearSegmentedColormap.from_list('t', ['#eef1f4', BLUE])]

fig, axes = plt.subplots(1, 2, figsize=(10, 3.1))
titles = ['counts (bag-of-words)', 'TF-IDF weights (each row has length 1)']
for ax, M, cmap, title, fmt in zip(axes, [C, T], cmaps, titles, ['{:.0f}', '{:.2f}']):
    ax.imshow(M, cmap=cmap, vmin=0, vmax=M.max(), aspect='auto')
    for i in range(M.shape[0]):
        for j in range(M.shape[1]):
            ax.text(j, i, fmt.format(M[i, j]), ha='center', va='center', fontsize=8.5,
                    color='black' if M[i, j] > 0 else '#9aa4ae')
    ax.set_xticks(range(len(vocab)))
    ax.set_xticklabels([f'{w}\ndf={d}' for w, d in zip(vocab, df)], fontsize=8)
    ax.set_yticks(range(3))
    ax.set_yticklabels(['d1', 'd2', 'd3'], fontsize=9)
    ax.set_title(title, fontsize=10)
    ax.tick_params(length=0)
    for s in ax.spines.values():
        s.set_visible(False)

fig.suptitle('"gato" is in every document: 2 counts in d3 become 0.64, while "peixe" (1 count, df=1) gets 0.72 in d1',
             fontsize=9.5)
plt.tight_layout()

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

plt.close()