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

# Left: the document-term count matrix of the 11,314 training posts of 20 Newsgroups.
# We draw its densest corner: 200 posts x the 2,000 most frequent words; each dot is
# a non-zero entry. Right: vocabulary size (columns) and share of zeros when bigrams
# and trigrams are added.
news = fetch_20newsgroups(subset='train', remove=('headers', 'footers', 'quotes'))

cv = CountVectorizer()
X = cv.fit_transform(news.data)
order = np.argsort(-np.asarray(X.sum(axis=0)).ravel())
corner = X[:200][:, order[:2000]]
zeros_all = 1 - X.nnz / (X.shape[0] * X.shape[1])
zeros_corner = 1 - corner.nnz / (corner.shape[0] * corner.shape[1])

ranges = [(1, 1), (1, 2), (1, 3)]
sizes, zeros, per_doc = [], [], []
for r in ranges:
    Xr = CountVectorizer(ngram_range=r).fit_transform(news.data)
    sizes.append(Xr.shape[1])
    zeros.append(1 - Xr.nnz / (Xr.shape[0] * Xr.shape[1]))
    per_doc.append(np.median(np.diff(Xr.indptr)))

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, 3.9), gridspec_kw={'width_ratios': [1.5, 1]})

rows, cols = corner.nonzero()
a1.scatter(cols, rows, s=0.6, color=BLUE, marker='s', linewidths=0, rasterized=True)
a1.set_xlim(0, 2000)
a1.set_ylim(200, 0)
a1.set_xlabel('the 2,000 most frequent words (of 101,631), most frequent on the left')
a1.set_ylabel('document')
a1.set_title(f'The densest corner of the matrix is still {zeros_corner:.1%} zeros;\n'
             f'the whole {X.shape[0]:,} x {X.shape[1]:,} matrix is {zeros_all:.2%} zeros', fontsize=9.5)

labels = ['words\n(1, 1)', '+ bigrams\n(1, 2)', '+ trigrams\n(1, 3)']
bars = a2.bar(labels, np.array(sizes) / 1e6, color=[BLUE, ORANGE, RED], width=0.6)
for b, s, z, m in zip(bars, sizes, zeros, per_doc):
    a2.text(b.get_x() + b.get_width() / 2, b.get_height() + 0.08,
            f'{s:,} columns\n{z:.2%} zeros\n{m:.0f} non-zero per doc',
            ha='center', va='bottom', fontsize=7.8)
a2.set_ylim(0, max(sizes) / 1e6 * 1.55)
a2.set_ylabel('vocabulary size (millions)')
a2.set_title(f'n-grams multiply the columns (x{sizes[2] / sizes[0]:.0f}),\nbut each post still fills only a few hundred', fontsize=9.5)
for s in ['top', 'right']:
    a2.spines[s].set_visible(False)

plt.tight_layout()

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

plt.close()