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 CountVectorizer, TfidfVectorizer
from sklearn.manifold import TSNE
from sklearn.neighbors import NearestNeighbors
from sklearn.preprocessing import normalize

# Posts from four newsgroups (those with at least 20 words once headers, signatures
# and quoted replies are removed), two representations: raw word counts and
# TF-IDF (English stop words removed). Each is reduced to 50 dimensions with a
# truncated SVD and drawn in 2-D with t-SNE. The color (the newsgroup) is never
# shown to any of the steps. The number in each title: the share of each post's
# 5 nearest neighbors (cosine, in the full sparse space) that come from its newsgroup.
cats = ['rec.autos', 'sci.med', 'sci.space', 'talk.politics.guns']
news = fetch_20newsgroups(subset='train', categories=cats,
                          remove=('headers', 'footers', 'quotes'), random_state=0)
keep = [i for i, t in enumerate(news.data) if len(t.split()) >= 20]
docs, y = [news.data[i] for i in keep], news.target[keep]


def agreement(V):
    nn = NearestNeighbors(n_neighbors=6, metric='cosine').fit(V)
    _, ind = nn.kneighbors(V)
    return (y[ind[:, 1:]] == y[:, None]).mean()


reps = {'raw counts': CountVectorizer().fit_transform(docs),
        'TF-IDF': TfidfVectorizer(stop_words='english').fit_transform(docs)}

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.4))
for ax, (name, V) in zip(axes, reps.items()):
    Z = normalize(TruncatedSVD(n_components=50, random_state=0).fit_transform(V))
    P = TSNE(2, perplexity=30, init='pca', random_state=0).fit_transform(Z)
    for c, col, lab in zip(range(4), [BLUE, ORANGE, GREEN, RED], cats):
        m = y == c
        ax.scatter(P[m, 0], P[m, 1], s=4, color=col, alpha=0.8, linewidths=0, label=lab, rasterized=True)
    ax.set_title(f'{name}: {agreement(V):.0%} of the 5 nearest posts\ncome from the same newsgroup',
                 fontsize=9.5)
    ax.set_xticks([])
    ax.set_yticks([])
axes[1].legend(fontsize=8, frameon=False, markerscale=3, loc='center left', bbox_to_anchor=(1.0, 0.5))
fig.suptitle(f'The same {len(docs):,} posts, same map-making (SVD to 50-d, then t-SNE); only the representation changes',
             fontsize=10)
plt.tight_layout()

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

plt.close()