import matplotlib.pyplot as plt
import numpy as np
from io import StringIO

# Topic mixtures of documents with 3 topics, drawn from Dirichlet(alpha) for three alphas
# (the same draws as the lesson's code: 2000 per alpha; 400 are plotted).
# Each point is one document's theta = (share of topic 0, 1, 2), placed inside the triangle
# whose corners are the "pure" documents.
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})

corners = np.array([[0, 0], [1, 0], [0.5, np.sqrt(3) / 2]])
rng = np.random.default_rng(1)
fig, axes = plt.subplots(1, 3, figsize=(9.5, 3.4))
for ax, alpha in zip(axes, [0.1, 1.0, 10.0]):
    theta = rng.dirichlet([alpha] * 3, size=2000)
    xy = theta[:400] @ corners
    colors = np.array([BLUE, ORANGE, GREEN])[theta[:400].argmax(axis=1)]
    ax.scatter(xy[:, 0], xy[:, 1], s=7, c=colors, alpha=0.7, linewidths=0)
    ax.plot(*np.vstack([corners, corners[:1]]).T, color=GRAY, lw=1)
    for (cx, cy), name, dy in zip(corners, ['topic 0', 'topic 1', 'topic 2'], [-0.09, -0.09, 0.05]):
        ax.text(cx, cy + dy, name, ha='center', fontsize=8.5)
    ax.set_title(f'α = {alpha:g}   (largest share, on average: {theta.max(axis=1).mean():.2f})',
                 fontsize=9)
    ax.set_xlim(-0.12, 1.12)
    ax.set_ylim(-0.15, 1.0)
    ax.set_aspect('equal')
    ax.axis('off')
fig.suptitle('Small α: each document sticks to one topic.  Large α: every document mixes all three.',
             fontsize=10)
plt.tight_layout()

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

plt.close()