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

# LDA's generative story with the lesson's numbers: two topics over six words,
# one document that is 80% "sports" and 20% "space"
BLUE, ORANGE, GRAY = '#58a6ff', '#f0883e', 'gray'
plt.rcParams.update({'text.color': GRAY, 'axes.labelcolor': GRAY, 'axes.edgecolor': GRAY,
                     'xtick.color': GRAY, 'ytick.color': GRAY})

vocab = ['team', 'goal', 'game', 'rocket', 'orbit', 'moon']
phi = np.array([[0.40, 0.30, 0.20, 0.05, 0.00, 0.05],
                [0.05, 0.00, 0.05, 0.40, 0.30, 0.20]])
theta = np.array([0.8, 0.2])
part = theta[:, None] * phi              # contribution of each topic to p(w|d)
p_word = part.sum(axis=0)

x = np.arange(len(vocab))
fig, axes = plt.subplots(1, 3, figsize=(10, 3.2), sharey=True)
for ax, row, color, title in [(axes[0], phi[0], BLUE, 'topic 0 ("sports"): $\\varphi_0$'),
                              (axes[1], phi[1], ORANGE, 'topic 1 ("space"): $\\varphi_1$')]:
    ax.bar(x, row, color=color, width=0.7)
    for xi, v in zip(x, row):
        ax.text(xi, v + 0.01, f'{v:.2f}', ha='center', fontsize=8)
    ax.set_title(title, fontsize=9.5)
ax = axes[2]
ax.bar(x, part[0], color=BLUE, width=0.7, label='0.8 × topic 0')
ax.bar(x, part[1], bottom=part[0], color=ORANGE, width=0.7, label='0.2 × topic 1')
for xi, v in zip(x, p_word):
    ax.text(xi, v + 0.01, f'{v:.2f}', ha='center', fontsize=8)
ax.set_title('the document: $p(w\\,|\\,d)=0.8\\,\\varphi_0+0.2\\,\\varphi_1$', fontsize=9.5)
ax.legend(fontsize=8, frameon=False, loc='upper right')
for ax in axes:
    ax.set_xticks(x, vocab, fontsize=8.5, rotation=30)
    ax.tick_params(axis='y', labelsize=8)
    ax.spines[['top', 'right']].set_visible(False)
axes[0].set_ylabel('probability of the word', fontsize=9)
axes[0].set_ylim(0, 0.47)
fig.suptitle('Every word of the document: draw a topic from θ = (0.8, 0.2), then a word from that topic',
             fontsize=10)
plt.tight_layout()

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

plt.close()