import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
from io import StringIO
from scipy.stats import loguniform, randint
from sklearn.datasets import fetch_openml
from sklearn.ensemble import HistGradientBoostingClassifier
from sklearn.experimental import enable_halving_search_cv  # noqa: F401
from sklearn.model_selection import (HalvingRandomSearchCV, ParameterSampler, StratifiedKFold,
                                     cross_val_score, train_test_split)

# The same 27 random configurations of a boosted model on kc1 (software defects), run
# through successive halving twice: once with trees as the budget (10 -> 270), once with
# rows (60 -> 1620). Each configuration is also scored at the full budget (270 trees, all
# rows), which is what halving tries to predict from its cheap early rounds.
kc1 = fetch_openml(data_id=1067, as_frame=True)
X, y = kc1.data, (kc1.target == 'true').astype(int)
X_dev, X_test, y_dev, y_test = train_test_split(X, y, test_size=0.2, stratify=y, random_state=0)
cv = StratifiedKFold(n_splits=5, shuffle=True, random_state=0)
space = {'learning_rate': loguniform(0.01, 0.5), 'max_leaf_nodes': randint(4, 64),
         'min_samples_leaf': randint(5, 200), 'l2_regularization': loguniform(1e-3, 10)}
base = HistGradientBoostingClassifier(max_iter=270, random_state=0)

runs = {}
for name, kw in [('trees', dict(resource='max_iter', min_resources=10, max_resources=270)),
                 ('rows', dict(min_resources=60))]:
    h = HalvingRandomSearchCV(base, space, n_candidates=27, factor=3, cv=cv, scoring='roc_auc',
                              random_state=0, **kw).fit(X_dev, y_dev)
    res = pd.DataFrame(h.cv_results_)
    res['key'] = [p['learning_rate'] for p in res['params']]          # identifies a candidate
    runs[name] = (res.groupby('key')['iter'].max(), h.best_params_['learning_rate'])

full = {}                                                  # every candidate at the full budget
for p in ParameterSampler(space, n_iter=27, random_state=0):     # the same 27 halving sampled
    full[p['learning_rate']] = cross_val_score(base.set_params(**p), X_dev, y_dev, cv=cv,
                                               scoring='roc_auc').mean()
full = pd.Series(full)
rank = full.rank(ascending=False)

fig, axes = plt.subplots(1, 2, figsize=(10, 4), sharey=True)
colors = {0: 'gray', 1: '#58a6ff', 2: '#f0883e'}
labels = {0: 'out after round 1', 1: 'out after round 2', 2: 'out after round 3'}
for ax, (name, (reached, winner)) in zip(axes, runs.items()):
    for r in (0, 1, 2):
        keys = reached[reached == r].index
        ax.scatter(keys, full[keys], color=colors[r], s=30, label=labels[r], zorder=2)
    ax.scatter([winner], [full[winner]], marker='*', s=220, color='#3fb950', zorder=3,
               label='winner (round 4)', edgecolor='white', lw=0.6)
    ax.set_xscale('log')
    ax.set_xlabel('learning rate of the configuration (log scale)', fontsize=8.5)
    what = '10 → 30 → 90 → 270 trees' if name == 'trees' else '60 → 180 → 540 → 1,620 rows'
    ax.set_title(f'budget = {name}: {what}\nwinner ranks {int(rank[winner])} of 27 at the full budget',
                 fontsize=9)
    ax.grid(alpha=0.25)
    ax.tick_params(labelsize=7.5)
axes[0].set_ylabel('CV ROC AUC at the full budget\n(270 trees, all rows)', fontsize=8.5)
axes[1].legend(fontsize=7.5, loc='lower left')
fig.suptitle('Successive halving on the same 27 configurations: what survives depends on the budget unit',
             fontsize=9.5)
plt.tight_layout()

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

plt.close()