from io import StringIO

import matplotlib.pyplot as plt
import numpy as np
from scipy.stats import ks_2samp

# The page's PSI experiment: the same draws, keeping the sample whose mean moved by 0.5 sd
rng = np.random.default_rng(0)
ref = rng.normal(0, 1, 5000)
for shift in [0.0, 0.1, 0.25, 0.5]:
    cur = rng.normal(shift, 1, 5000)

edges = np.quantile(ref, np.linspace(0, 1, 11))
edges[0], edges[-1] = -np.inf, np.inf
e = np.histogram(ref, edges)[0] / len(ref)
a = np.histogram(cur, edges)[0] / len(cur)
terms = (a - e) * np.log(a / e)
psi = terms.sum()

fig, (a1, a2) = plt.subplots(1, 2, figsize=(9.5, 3.6))
idx = np.arange(10)
a1.bar(idx - 0.2, e * 100, width=0.4, color='gray', alpha=0.6, label='reference (training)')
a1.bar(idx + 0.2, a * 100, width=0.4, color='#f0883e', label='current (mean + 0.5 sd)')
for i in idx:
    a1.text(i + 0.2, a[i] * 100 + 0.4, f'{terms[i]:.3f}', ha='center', fontsize=6.5, color='gray', rotation=90)
a1.set_xticks(idx, [f'D{i + 1}' for i in idx], fontsize=8)
a1.set_xlabel('decile bins of the reference')
a1.set_ylabel('% of rows')
a1.set_ylim(0, 25)
a1.set_title(f'PSI = sum of the per-bin terms (gray) = {psi:.2f}', fontsize=9.5, loc='left')
a1.legend(fontsize=7.5, loc='upper left', frameon=False)

xs = np.linspace(-4, 4.5, 800)
F_ref = np.searchsorted(np.sort(ref), xs, side='right') / len(ref)
F_cur = np.searchsorted(np.sort(cur), xs, side='right') / len(cur)
res = ks_2samp(ref, cur)
i = np.argmax(np.abs(F_ref - F_cur))
a2.plot(xs, F_ref, color='gray', lw=2, label='reference')
a2.plot(xs, F_cur, color='#f0883e', lw=2, label='current')
a2.vlines(xs[i], F_cur[i], F_ref[i], color='#58a6ff', lw=2.5)
a2.annotate(f'D = {res.statistic:.2f}', xy=(xs[i], (F_cur[i] + F_ref[i]) / 2), xytext=(xs[i] + 0.8, 0.35),
            fontsize=9, color='#58a6ff', arrowprops=dict(arrowstyle='->', color='#58a6ff'))
a2.set_xlabel('x')
a2.set_ylabel('share of rows ≤ x')
a2.set_title(f'KS: largest gap between the two CDFs (p = {res.pvalue:.0e})', fontsize=9.5, loc='left')
a2.legend(fontsize=7.5, loc='upper left', frameon=False)
for ax in (a1, a2):
    ax.tick_params(labelsize=8)
    ax.grid(alpha=0.2)
plt.tight_layout()

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