from io import StringIO

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

# The page's simulation of the three shifts (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)
model = HistGradientBoostingClassifier(max_depth=3, random_state=0).fit(x_ref.reshape(-1, 1), y_ref)
pool_x, pool_y = draw(100_000, 0, rng)
pos, neg = pool_x[pool_y == 1][:1000], pool_x[pool_y == 0][:4000]
scenarios = {'no change': draw(5000, 0, rng),
             'covariate shift': draw(5000, 2.0, rng),
             'concept drift': draw(5000, 0, rng, move=1.0),
             'prior shift': (np.r_[pos, neg], np.r_[np.ones(1000), np.zeros(4000)].astype(int))}
prior_ref = pool_y.mean()

grid = np.linspace(-3.5, 5.5, 400)
model_curve = model.predict_proba(grid.reshape(-1, 1))[:, 1]
bins = np.linspace(-3.5, 5.5, 46)

fig, axes = plt.subplots(1, 4, figsize=(10, 3.3), sharey=True)
for ax, (name, (x, y)) in zip(axes, scenarios.items()):
    h_ref = np.histogram(x_ref, bins, density=True)[0]
    h_cur = np.histogram(x, bins, density=True)[0]
    scale = 0.9 / max(h_ref.max(), h_cur.max())
    ax.stairs(h_ref * scale, bins, fill=True, color='gray', alpha=0.25, label='x in training')
    ax.stairs(h_cur * scale, bins, color='#f0883e', lw=1.3, label='x now')
    if name == 'concept drift':
        truth = true_p(grid, 1.0)
    elif name == 'prior shift':
        p = true_p(grid)
        w1, w0 = 0.2 / prior_ref, 0.8 / (1 - prior_ref)
        truth = p * w1 / (p * w1 + (1 - p) * w0)
    else:
        truth = true_p(grid)
    ax.plot(grid, truth, color='#3fb950', lw=2, label='true P(y=1 | x) now')
    ax.plot(grid, model_curve, color='#58a6ff', lw=2, ls='--', label='model (trained once)')
    acc = np.mean((model.predict_proba(x.reshape(-1, 1))[:, 1] > 0.5) == y)
    ax.set_title(f'{name}\nPSI {psi(x_ref, x):.2f} · accuracy {acc:.2f}', fontsize=9)
    ax.set_xlabel('x', fontsize=8)
    ax.set_ylim(0, 1.05)
    ax.tick_params(labelsize=7.5)
axes[0].set_ylabel('probability  (histograms scaled)', fontsize=8)
handles, labels = axes[0].get_legend_handles_labels()
fig.legend(handles, labels, loc='lower center', ncol=4, fontsize=8, frameon=False)
plt.tight_layout(rect=(0, 0.08, 1, 1))

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