import matplotlib.pyplot as plt
import numpy as np
from io import StringIO
from sklearn.datasets import fetch_openml
from sklearn.ensemble import HistGradientBoostingClassifier
from sklearn.model_selection import train_test_split

# The lab's development set (bank marketing, call duration removed). Each model is allowed
# 1,000 trees; HistGradientBoosting records the log loss on its inner 10% validation slice
# after every tree. Open circle: the minimum; dot: where early stopping (patience 30) stops.
# The fitted model keeps every tree up to the dot, including the 30 that did not help.
bank = fetch_openml(data_id=1461, as_frame=True)
X = bank.data.drop(columns='V12')                   # V12 = call duration: known only after the call
y = (bank.target == '2').astype(int)
X_dev, X_test, y_dev, y_test = train_test_split(X, y, test_size=0.2, stratify=y, random_state=42)

fig, ax = plt.subplots(figsize=(7.5, 3.6))
for eta, c in [(0.5, '#ff7b72'), (0.1, '#58a6ff'), (0.03, '#3fb950')]:
    m = HistGradientBoostingClassifier(learning_rate=eta, max_iter=1000, early_stopping=True,
                                       n_iter_no_change=1000, random_state=42).fit(X_dev, y_dev)
    loss = -m.validation_score_[1:]                  # entry 0 is the constant model F0
    trees = np.arange(1, len(loss) + 1)
    s = m.validation_score_                          # scikit-learn's own stopping rule, patience 30
    stop = next((k - 1 for k in range(32, len(s) + 1)
                 if not any(v > s[k - 31] + 1e-7 for v in s[k - 30:k])), len(s) - 1)
    b = int(loss.argmin())
    ax.plot(trees, loss, color=c, lw=1.6,
            label=f'η = {eta}: best {loss[b]:.3f} at {b + 1} trees · stops at {stop} trees, {loss[stop - 1]:.3f}')
    ax.scatter([b + 1], [loss[b]], s=40, zorder=4, facecolors='none', edgecolors=c, lw=1.4)
    ax.scatter([stop], [loss[stop - 1]], color=c, s=36, zorder=4, edgecolor='white')
ax.set_xscale('log')
ax.set_ylim(0.275, 0.35)
ax.set_xlabel('number of trees (log scale)', fontsize=8.5)
ax.set_ylabel('log loss on the inner validation 10%', fontsize=8.5)
ax.set_title('Bank marketing: early stopping picks the number of trees for each learning rate', fontsize=9.5)
ax.legend(fontsize=7.5, loc='upper right', title='open circle = minimum · dot = where early stopping stops',
          title_fontsize=7.5)
ax.tick_params(labelsize=7.5)
ax.grid(alpha=0.25)

plt.tight_layout()

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

plt.close()