import sys
from io import StringIO

import matplotlib.pyplot as plt
import numpy as np
from joblib import Parallel, delayed
from scipy.stats import loguniform, randint
from sklearn.datasets import fetch_openml
from sklearn.ensemble import HistGradientBoostingClassifier
from sklearn.metrics import roc_auc_score
from sklearn.model_selection import (ParameterSampler, StratifiedKFold, cross_val_score,
                                     train_test_split)

# Random search on kc1 (software defects). A development sample of 1,687 rows (or only 300)
# is drawn 10 times; the rest of the table is the test set. For each draw, 150 random
# configurations are scored by 5-fold CV on the development rows and, for reference only,
# on the test rows. A search of n trials keeps the configuration with the best CV score;
# we record its CV score and its test score, averaged over 300 random search orders.
kc1 = fetch_openml(data_id=1067, as_frame=True)
X, y = kc1.data, (kc1.target == 'true').astype(int)
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)}
N_POOL, SPLITS, ORDERS = 150, 10, 300


def evaluate(params, X_dev, y_dev, X_test, y_test, cv):
    model = HistGradientBoostingClassifier(random_state=0, **params)
    cv_auc = cross_val_score(model, X_dev, y_dev, cv=cv, scoring='roc_auc').mean()
    test_auc = roc_auc_score(y_test, model.fit(X_dev, y_dev).predict_proba(X_test)[:, 1])
    return cv_auc, test_auc


def curves(n_dev):
    cv_curves, test_curves = [], []
    rng = np.random.default_rng(0)
    for split in range(SPLITS):
        X_dev, X_test, y_dev, y_test = train_test_split(X, y, train_size=n_dev, stratify=y,
                                                        random_state=split)
        cv = StratifiedKFold(n_splits=5, shuffle=True, random_state=split)
        pool = list(ParameterSampler(space, n_iter=N_POOL, random_state=split))
        scores = np.array(Parallel(n_jobs=4)(delayed(evaluate)(p, X_dev, y_dev, X_test, y_test, cv)
                                             for p in pool))
        for _ in range(ORDERS):                           # many random search orders
            order = rng.permutation(N_POOL)
            cv_s, te_s = scores[order, 0], scores[order, 1]
            best = np.maximum.accumulate(cv_s)
            chosen = np.array([np.flatnonzero(cv_s[:n] == best[n - 1])[0] for n in range(1, N_POOL + 1)])
            cv_curves.append(cv_s[chosen])
            test_curves.append(te_s[chosen])
    return np.mean(cv_curves, axis=0), np.mean(test_curves, axis=0)


n = np.arange(1, N_POOL + 1)
fig, axes = plt.subplots(1, 2, figsize=(10, 3.9))
for ax, n_dev in zip(axes, [1687, 300]):
    cv_mean, test_mean = curves(n_dev)
    for k in (1, 10, 30, 150):
        print(f"n_dev={n_dev} trials={k:3d}  CV {cv_mean[k - 1]:.4f}  test {test_mean[k - 1]:.4f}",
              file=sys.stderr)
    ax.plot(n, cv_mean, color='#58a6ff', lw=1.8, label='CV score of the chosen configuration')
    ax.plot(n, test_mean, color='#f0883e', lw=1.8, label='test score of the same configuration')
    for k in (1, 10, 150):
        hi_cv = cv_mean[k - 1] >= test_mean[k - 1]              # label the higher line above
        ax.annotate(f'{cv_mean[k - 1]:.3f}', (k, cv_mean[k - 1]), textcoords='offset points',
                    xytext=(0, 6 if hi_cv else -13), fontsize=7.5, color='#58a6ff', ha='center')
        ax.annotate(f'{test_mean[k - 1]:.3f}', (k, test_mean[k - 1]), textcoords='offset points',
                    xytext=(0, -13 if hi_cv else 6), fontsize=7.5, color='#f0883e', ha='center')
    ax.set_xscale('log')
    ax.set_xlim(0.7, 220)
    lo = min(cv_mean.min(), test_mean.min())
    hi = max(cv_mean.max(), test_mean.max())
    ax.set_ylim(lo - 0.15 * (hi - lo), hi + 0.15 * (hi - lo))   # room for the labels
    ax.set_xlabel('number of configurations tried (log scale)', fontsize=8.5)
    ax.set_title(f'{n_dev:,} development rows (test: the other {len(y) - n_dev:,})', fontsize=9.5)
    ax.grid(alpha=0.25)
    ax.tick_params(labelsize=7.5)
axes[0].set_ylabel('ROC AUC', fontsize=8.5)
axes[0].legend(fontsize=7.5, loc='lower right')
fig.suptitle('Random search on kc1: what a longer search buys, as reported and as delivered', fontsize=9.5)
fig.text(0.01, 0.01, f'mean over {SPLITS} development/test splits × {ORDERS} random orders of {N_POOL} configurations',
         fontsize=7, color='gray')
plt.tight_layout(rect=(0, 0.04, 1, 1))

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

plt.close()