import matplotlib.pyplot as plt
import numpy as np
from io import StringIO
from sklearn.datasets import fetch_20newsgroups
from sklearn.decomposition import NMF, LatentDirichletAllocation
from sklearn.feature_extraction.text import CountVectorizer, TfidfVectorizer
from sklearn.metrics import adjusted_rand_score

# The lesson's corpus (2155 posts from 4 forums). Cut every post to its first n words,
# fit LDA (counts) and NMF (TF-IDF) with 4 topics, send each post to its largest topic and
# measure the agreement with the forums (ARI). Three seeds per point.
BLUE, ORANGE, GRAY = '#58a6ff', '#f0883e', 'gray'
plt.rcParams.update({'text.color': GRAY, 'axes.labelcolor': GRAY, 'axes.edgecolor': GRAY,
                     'xtick.color': GRAY, 'ytick.color': GRAY})

forums = ['comp.graphics', 'rec.sport.hockey', 'sci.space', 'soc.religion.christian']
ng = fetch_20newsgroups(subset='train', categories=forums, random_state=0,
                        remove=('headers', 'footers', 'quotes'))
long_enough = np.array([len(text.split()) >= 20 for text in ng.data])
docs = [text for text, keep in zip(ng.data, long_enough) if keep]
y = ng.target[long_enough]
settings = dict(stop_words='english', min_df=2, max_df=0.5, token_pattern=r'(?u)\b[a-zA-Z]{3,}\b')

lengths = [5, 10, 20, 40, 80, 160, None]
res = {'LDA': [], 'NMF': []}
for n in lengths:
    cut = [' '.join(t.split()[:n]) for t in docs] if n else docs
    X = CountVectorizer(**settings).fit_transform(cut)
    T = TfidfVectorizer(**settings).fit_transform(cut)
    res['LDA'].append([adjusted_rand_score(y, LatentDirichletAllocation(n_components=4, random_state=s)
                                           .fit_transform(X).argmax(axis=1)) for s in range(3)])
    res['NMF'].append([adjusted_rand_score(y, NMF(n_components=4, init='nndsvda', max_iter=500,
                                                  random_state=s).fit_transform(T).argmax(axis=1))
                       for s in range(3)])
median_len = int(np.median([len(t.split()) for t in docs]))

x = np.arange(len(lengths))
fig, ax = plt.subplots(figsize=(7.5, 3.6))
for name, color in [('NMF', BLUE), ('LDA', ORANGE)]:
    r = np.array(res[name])
    ax.fill_between(x, r.min(axis=1), r.max(axis=1), color=color, alpha=0.2, lw=0)
    ax.plot(x, r.mean(axis=1), 'o-', color=color, label=f'{name} (mean of 3 seeds, band = min–max)')
    for xi, v in zip(x, r.mean(axis=1)):
        if xi in (2, len(x) - 1):
            ax.annotate(f'{v:.2f}', (xi, v), textcoords='offset points', xytext=(8, -4), fontsize=8.5,
                        color=color)
ax.set_xticks(x, [str(n) for n in lengths[:-1]] + [f'whole post\n(median {median_len})'], fontsize=8.5)
ax.set_xlabel('words kept from the start of each post', fontsize=9)
ax.set_ylabel('agreement with the forums (ARI)', fontsize=9)
ax.tick_params(axis='y', labelsize=8)
ax.set_ylim(-0.03, 1)
ax.spines[['top', 'right']].set_visible(False)
ax.legend(fontsize=8, frameon=False, loc='upper left')
ax.set_title('Word-count topic models need words: below ~40 words per post, the topics dissolve',
             fontsize=10)
plt.tight_layout()

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

plt.close()