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 TfidfVectorizer

# idf as a function of document frequency, for the n = 11,314 training posts of
# 20 Newsgroups: scikit-learn's smoothed formula ln((1+n)/(1+df)) + 1 and the
# textbook ln(n/df). A few words of the corpus are placed on the curve.
news = fetch_20newsgroups(subset='train', remove=('headers', 'footers', 'quotes'))
tfidf = TfidfVectorizer().fit(news.data)
n = len(news.data)
X = tfidf.transform(news.data)
df_all = np.bincount(X.indices, minlength=X.shape[1])
vocab = {w: i for i, w in enumerate(tfidf.get_feature_names_out())}

df = np.logspace(0, np.log10(n), 300)
smooth = np.log((1 + n) / (1 + df)) + 1
classic = np.log(n / df)

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, ax = plt.subplots(figsize=(8, 4.2))
ax.semilogx(df, smooth, color=BLUE, lw=2.5, label='scikit-learn: $\\ln\\frac{1+n}{1+df} + 1$')
ax.semilogx(df, classic, '--', color='gray', lw=1.5, label='textbook: $\\ln\\frac{n}{df}$')
words = ['the', 'people', 'car', 'hockey', 'odometer', 'filioque']
for w in words:
    d = df_all[vocab[w]]
    v = tfidf.idf_[vocab[w]]
    ax.plot(d, v, 'o', color=ORANGE, ms=5, zorder=3)
    ax.annotate(f'{w}\n(df={d:,})', (d, v), xytext=(6, 6), textcoords='offset points', fontsize=7.8,
                color=ORANGE)
ax.axhline(1, color=GREEN, lw=0.8, ls=':')
ax.text(1.2, 1.25, 'a word in every post still weighs 1 (textbook: 0)', fontsize=8, color=GREEN)
ax.set_xlabel(f'df: number of posts containing the word (n = {n:,})')
ax.set_ylabel('idf')
ax.set_ylim(-0.3, 11)
ax.set_title(f'idf falls from {smooth[0]:.1f} for a word in one post to 1 for a word in all of them',
             fontsize=10)
ax.legend(fontsize=8.5, frameon=False, loc='upper right')
for s in ['top', 'right']:
    ax.spines[s].set_visible(False)
plt.tight_layout()

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

plt.close()