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

plt.rcParams.update({'text.color': 'gray', 'axes.labelcolor': 'gray', 'axes.edgecolor': 'gray',
                     'xtick.color': 'gray', 'ytick.color': 'gray'})

# The lesson's simulation: 10,000 students, ability ~ N(70, 8), each exam = ability + luck.
# Same seed and same order of draws as the code in the lesson.
rng = np.random.default_rng(0)
n = 10_000
ability = rng.normal(70, 8, n)
exams = {}
for luck in [2, 6, 16]:
    exams[luck] = (ability + rng.normal(0, luck, n), ability + rng.normal(0, luck, n))

show = np.random.default_rng(1).choice(n, 900, replace=False)      # points drawn, for a light SVG
fig, axes = plt.subplots(1, 2, figsize=(9.4, 4.2))
for ax, luck in zip(axes, [6, 16]):
    e1, e2 = exams[luck]
    r = 64 / (64 + luck ** 2)
    top, bottom = e1 >= np.quantile(e1, 0.9), e1 <= np.quantile(e1, 0.1)
    rest = ~top & ~bottom
    ax.scatter(e1[show][rest[show]], e2[show][rest[show]], s=5, color='gray', alpha=0.35, lw=0)
    ax.scatter(e1[show][top[show]], e2[show][top[show]], s=7, color='#f0883e', alpha=0.7, lw=0,
               label='top 10% on exam 1')
    ax.scatter(e1[show][bottom[show]], e2[show][bottom[show]], s=7, color='#58a6ff', alpha=0.7, lw=0,
               label='bottom 10% on exam 1')
    lim = (15, 125)
    ax.plot(lim, lim, '--', color='gray', lw=1, label='exam 2 = exam 1')
    m = e1.mean()
    xs = np.array(lim)
    ax.plot(xs, m + r * (xs - m), color='#3fb950', lw=2, label=f'expected exam 2 (slope r = {r:.2f})')
    for grp, color in [(top, '#f0883e'), (bottom, '#58a6ff')]:
        a, b = e1[grp].mean(), e2[grp].mean()
        ax.plot([a], [b], 'o', color=color, ms=8, mec='white', mew=1)
        label = f'{"top" if grp is top else "bottom"} 10%:\nexam 1 {a:.1f}\nexam 2 {b:.1f}'
        ax.annotate(label, xy=(a, b), xytext=(93, 22) if grp is top else (18, 98),
                    fontsize=8, color=color, arrowprops=dict(arrowstyle='->', color=color))
    ax.set_xlim(lim)
    ax.set_ylim(lim)
    ax.set_aspect('equal')
    ax.set_xlabel('exam 1 score')
    ax.set_title(f'luck sd {luck}:  reliability r = {r:.2f}', fontsize=10)
axes[0].set_ylabel('exam 2 score (same students)')
handles, labels = axes[0].get_legend_handles_labels()
fig.legend(handles, [l.replace(' (slope r = 0.64)', ' (slope r)') for l in labels], loc='lower center',
           ncol=4, fontsize=8, frameon=False, markerscale=2)
plt.tight_layout(rect=(0, 0.06, 1, 1))

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

plt.close()