import matplotlib.pyplot as plt
import numpy as np
from io import StringIO
from sklearn.datasets import fetch_20newsgroups
from sklearn.feature_extraction.text import CountVectorizer

# The 15 most frequent words of the 11,314 training posts of 20 Newsgroups,
# counted as a bag-of-words would count them: with scikit-learn's default
# tokenizer (left) and with its English stop-word list removed (right).
news = fetch_20newsgroups(subset='train', remove=('headers', 'footers', 'quotes'))

panels = []
for params in [{}, {'stop_words': 'english'}]:
    cv = CountVectorizer(**params)
    X = cv.fit_transform(news.data)
    freq = np.asarray(X.sum(axis=0)).ravel()
    top = np.argsort(-freq)[:15]
    panels.append((cv.get_feature_names_out()[top], freq[top], freq.sum(),
                   (X[:, top[0]] > 0).sum()))

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=(10, 4.3))
titles = ['default tokenizer', "stop_words='english'"]
for ax, (words, freq, total, df_top), color, title in zip(axes, panels, [ORANGE, BLUE], titles):
    y = np.arange(len(words))[::-1]
    ax.barh(y, freq / 1000, color=color, height=0.7)
    ax.set_yticks(y)
    ax.set_yticklabels(words, fontsize=8.5)
    share = freq.sum() / total
    ax.set_title(f'{title}: these 15 words are {share:.0%} of all tokens\n'
                 f'"{words[0]}" is in {df_top:,} of the {len(news.data):,} posts', fontsize=9.5)
    ax.set_xlabel('occurrences in the corpus (thousands)')
    for s in ['top', 'right']:
        ax.spines[s].set_visible(False)
axes[1].annotate('binary files pasted as text\nin a handful of posts', (freq[0] / 1000, y[0]),
                 xytext=(0.45, 0.62), textcoords='axes fraction', fontsize=8, color=RED,
                 arrowprops=dict(arrowstyle='->', color=RED, lw=1))

fig.suptitle('Raw counts are led by the words that say least about a post', fontsize=10)
plt.tight_layout()

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

plt.close()