import os, sys
import matplotlib.pyplot as plt
import numpy as np
from io import StringIO
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from _textstats import stats, STYLE

# Zipf's law: how many times each word occurs, against its rank (1 = most frequent),
# on a log-log scale. A few words are everywhere; most words occur once.
s = stats()
counts, vocab = s['counts'], s['vocab']
freq = np.sort(counts)[::-1]
hapax = (counts == 1).mean()
top10 = freq[:10].sum() / freq.sum()
GREEN, RED = '#3fb950', '#ff7b72'
plt.rcParams.update(STYLE)
fig, ax = plt.subplots(figsize=(6.6, 4.6))

r = np.arange(1, len(freq) + 1)
keep = np.unique(np.round(np.logspace(0, np.log10(len(freq)), 400)).astype(int)) - 1
ax.loglog(r[keep], freq[keep], color=GREEN, lw=3, drawstyle='steps-post')
first_one = np.argmax(freq == 1) + 1
ax.axhline(1, color=RED, lw=1, ls='--')
ax.axvspan(first_one, len(freq), color=RED, alpha=0.12)
ax.annotate(f'{hapax:.0%} of the {len(freq):,} words occur exactly once', xy=(first_one * 1.8, 1),
            xytext=(1.5e2, 0.45), color=RED, fontsize=10, ha='center', va='center',
            arrowprops=dict(arrowstyle='->', color=RED, lw=1))

LABEL = {'the': (25, 1.0e5), 'of': (70, 3.2e4), 'is': (200, 1.1e4), 'car': (2.5e3, 4.0e3), 'hockey': (6e3, 1.3e3)}
for w, (tx, ty) in LABEL.items():
    j = np.flatnonzero(freq == counts[vocab[w]])[0] + 1
    ax.plot(j, freq[j - 1], 'o', color=GREEN, ms=6, zorder=5)
    ax.annotate(f'"{w}": rank {j:,}, {freq[j - 1]:,}×', (j, freq[j - 1]), xytext=(tx, ty),
                fontsize=9.5, color='gray', va='center', arrowprops=dict(arrowstyle='-', color='gray', lw=0.6))
ax.text(1.2, 6, f'the 10 most frequent words\nare {top10:.0%} of all tokens', fontsize=10, color=GREEN)

ax.set_xlabel('rank of the word, 1 = most frequent (log scale)')
ax.set_ylabel('occurrences in the corpus (log scale)')
ax.set_title('Zipf: a few words are everywhere,\nmost words are rare', fontsize=11.5)
ax.set_xlim(0.8, len(freq) * 1.5)
ax.set_ylim(0.25, 3e5)
plt.tight_layout()

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