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.model_selection import train_test_split

# How many topics? On the lesson's corpus (2155 posts, 4 forums), for k = 2..20:
#  left:  LDA's perplexity on 20% of the posts held out from fitting (lower = "better")
#  right: coherence (mean NPMI of each topic's top-10 words, co-occurrence in the posts),
#         LDA and NMF, three seeds each
BLUE, ORANGE, GREEN, GRAY = '#58a6ff', '#f0883e', '#3fb950', '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'))
docs = [text for text in ng.data if len(text.split()) >= 20]
pattern = r'(?u)\b[a-zA-Z]{3,}\b'
settings = dict(stop_words='english', min_df=5, max_df=0.5, token_pattern=pattern)
cv = CountVectorizer(**settings)
X = cv.fit_transform(docs)
T = TfidfVectorizer(**settings).fit_transform(docs)
vocab = cv.get_feature_names_out()
X_fit, X_held = train_test_split(X, test_size=0.2, random_state=0)

# reference for coherence: which post contains which word
B = (X > 0).astype(np.float64).tocsc()
D = B.shape[0]


def coherence(components, n=10):
    scores = []
    for row in components:
        cols = B[:, row.argsort()[::-1][:n]].toarray()
        co = cols.T @ cols / D                     # co[i, j] = P(word i and word j in a post)
        p = np.diag(co)
        for i in range(n):
            for j in range(i + 1, n):
                scores.append(-1.0 if co[i, j] == 0 else
                              np.log(co[i, j] / (p[i] * p[j])) / -np.log(co[i, j]))
    return np.mean(scores)


ks = [2, 3, 4, 5, 6, 8, 10, 12, 15, 20]
perp, coh_lda, coh_nmf = [], [], []
for k in ks:
    p_row, c_row = [], []
    for seed in range(3):
        lda = LatentDirichletAllocation(n_components=k, random_state=seed).fit(X_fit)
        p_row.append(lda.perplexity(X_held))
        c_row.append(coherence(lda.components_))
    perp.append(p_row)
    coh_lda.append(c_row)
    coh_nmf.append([coherence(NMF(n_components=k, init='nndsvda', max_iter=500, random_state=seed)
                              .fit(T).components_) for seed in range(3)])
perp, coh_lda, coh_nmf = map(np.array, (perp, coh_lda, coh_nmf))

fig, (a1, a2) = plt.subplots(1, 2, figsize=(10, 3.6))
a1.fill_between(ks, perp.min(1), perp.max(1), color=ORANGE, alpha=0.2, lw=0)
a1.plot(ks, perp.mean(1), 'o-', color=ORANGE, label='LDA')
a1.set_title('held-out perplexity only rises: it would pick k = 2', fontsize=9.5)
a1.set_ylabel('perplexity on held-out posts', fontsize=9)
for arr, color, name in [(coh_nmf, BLUE, 'NMF'), (coh_lda, ORANGE, 'LDA')]:
    a2.fill_between(ks, arr.min(1), arr.max(1), color=color, alpha=0.2, lw=0)
    a2.plot(ks, arr.mean(1), 'o-', color=color, label=f'{name} (mean of 3 seeds)')
best = ks[int(np.argmax(coh_lda.mean(1)))]
a2.set_title(f'coherence: LDA peaks at k = {best}; NMF is higher at every k', fontsize=9.5)
a2.set_ylabel('mean NPMI of top-10 words', fontsize=9)
for ax in (a1, a2):
    ax.axvline(4, color=GREEN, lw=1, ls='--')
    ax.text(4.2, ax.get_ylim()[1], '4 forums', color=GREEN, fontsize=8, va='top')
    ax.set_xlabel('number of topics k', fontsize=9)
    ax.set_xticks(ks)
    ax.tick_params(labelsize=8)
    ax.spines[['top', 'right']].set_visible(False)
    ax.legend(fontsize=8, frameon=False, loc='lower right')
plt.tight_layout()

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

plt.close()