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

# Heaps' law: vocabulary size V after reading N tokens, on a log-log scale.
# A power law V = K N^beta is a straight line there; its slope is beta.
s = stats()
curve, K, beta = s['curve'], s['K'], s['beta']
BLUE, ORANGE = '#58a6ff', '#f0883e'
plt.rcParams.update(STYLE)
fig, ax = plt.subplots(figsize=(6.6, 4.6))

step = max(1, len(curve) // 400)
ax.loglog(curve[::step, 0], curve[::step, 1], color=BLUE, lw=3, label='measured on 20 Newsgroups')
N = np.logspace(3, np.log10(curve[-1, 0]), 50)
ax.loglog(N, K * N ** beta, '--', color=ORANGE, lw=1.8, label=f'fit  $V = {K:.1f}\\,N^{{{beta:.2f}}}$')
ax.loglog(N, N, ':', color='gray', lw=1.2, label='$V = N$: every token a new word')

# the three readings printed on the page
SPOT = {10_000: (2.4e4, 1.3e3), 100_000: (2.4e5, 7.5e3), 1_000_000: (1.4e5, 2.6e5)}
for target, (tx, ty) in SPOT.items():
    i = np.searchsorted(curve[:, 0], target)
    n, v = curve[i]
    ax.plot(n, v, 'o', color=BLUE, ms=7, zorder=5)
    ax.annotate(f'{n:,.0f} tokens → {v:,.0f} words', (n, v), xytext=(tx, ty), fontsize=9.5, color=BLUE,
                va='center', arrowprops=dict(arrowstyle='-', color=BLUE, lw=0.7))

# what the slope means
ax.text(2.0e6, 420, f'slope {beta:.2f}: ×10 text → ×{10 ** beta:.1f} words', color=ORANGE, fontsize=10.5, ha='right')

ax.set_xlabel('tokens read so far, N (log scale)')
ax.set_ylabel('distinct words seen, V (log scale)')
ax.set_title(f'Heaps: {curve[-1, 0]:,.0f} tokens → {curve[-1, 1]:,.0f} distinct words,\n'
             f'and the curve never flattens', fontsize=11.5)
ax.legend(fontsize=9.5, frameon=False, loc='upper left')
ax.set_xlim(1e3, curve[-1, 0] * 1.4)
ax.set_ylim(300, 3e6)
ax.set_xlim(1e3, curve[-1, 0] * 2.2)
plt.tight_layout()

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