from io import StringIO

import matplotlib.pyplot as plt
import numpy as np
from sklearn.ensemble import HistGradientBoostingClassifier

# The page's "degrade, detect, retrain" simulation (same seeds, same draws)


def true_p(x, move=0.0):
    z = x - move
    return 1 / (1 + np.exp(-(3 * z - z ** 2)))


def draw(n, mean, rng, move=0.0):
    x = rng.normal(mean, 1, n)
    return x, (rng.random(n) < true_p(x, move)).astype(int)


def psi(ref, cur, bins=10, eps=1e-4):
    edges = np.unique(np.quantile(ref, np.linspace(0, 1, bins + 1)))
    edges[0], edges[-1] = -np.inf, np.inf
    e = np.clip(np.histogram(ref, edges)[0] / len(ref), eps, None)
    a = np.clip(np.histogram(cur, edges)[0] / len(cur), eps, None)
    return float(np.sum((a - e) * np.log(a / e)))


rng = np.random.default_rng(2)
x_ref, y_ref = draw(5000, 0, rng)

rng = np.random.default_rng(3)
new_model = lambda: HistGradientBoostingClassifier(max_depth=3, random_state=0)
static = new_model().fit(x_ref.reshape(-1, 1), y_ref)
live, live_ref = static, x_ref
months, psi_static, psi_live, acc_s, acc_l, best, alerts = [], [], [], [], [], [], []
for month in range(13):
    x, y = draw(1000, 0.3 * month, rng)
    X_m = x.reshape(-1, 1)
    d = psi(live_ref, x)
    months.append(month)
    psi_static.append(psi(x_ref, x))
    psi_live.append(d)
    acc_s.append(np.mean(static.predict(X_m) == y))
    acc_l.append(np.mean(live.predict(X_m) == y))
    best.append(np.mean((true_p(x) > 0.5) == y))
    if d > 0.25:
        alerts.append(month)
        live, live_ref = new_model().fit(X_m, y), x

fig, (a1, a2) = plt.subplots(2, 1, figsize=(8, 5.2), sharex=True)
a1.plot(months, psi_static, 'o-', color='#f0883e', ms=4, lw=2, label='PSI vs the static model\'s training data')
a1.plot(months, psi_live, 'o-', color='#58a6ff', ms=4, lw=2, label='PSI vs the served model\'s training data')
a1.axhline(0.25, color='#ff7b72', lw=1, ls='--')
a1.text(0.2, 0.27, 'alert 0.25', color='#ff7b72', fontsize=8, va='bottom')
for m in alerts:
    a1.axvline(m, color='#ff7b72', alpha=0.15, lw=6)
a1.set_yscale('symlog', linthresh=0.3)
a1.set_yticks([0, 0.1, 0.25, 1, 4], ['0', '0.1', '0.25', '1', '4'])
a1.set_ylabel('PSI of x')
a1.set_ylim(0, 8)
a1.legend(fontsize=7.5, loc='upper left', frameon=False)
a1.set_title('The population moves 0.3 sd per month; red bands = alert, then retrain on that month',
             fontsize=9.5, loc='left')

a2.plot(months, best, color='gray', lw=5, alpha=0.35, label='best possible')
a2.plot(months, acc_l, 'o-', color='#3fb950', ms=4, lw=2, label='retrained on alert')
a2.plot(months, acc_s, 'o-', color='#f0883e', ms=4, lw=2, label='static model')
for m in alerts:
    a2.axvline(m, color='#ff7b72', alpha=0.15, lw=6)
a2.axhline(0.5, color='gray', lw=0.8, ls=':')
a2.text(0.1, 0.51, 'coin flip', color='gray', fontsize=7.5)
a2.annotate(f'{acc_s[-1]:.3f}', xy=(12, acc_s[-1]), xytext=(10.6, 0.3), fontsize=8, color='#f0883e')
a2.set_ylabel('accuracy')
a2.set_ylim(0.2, 0.95)
a2.set_xlabel('month')
a2.set_xticks(months)
a2.legend(fontsize=7.5, loc='lower 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()