# Needs optuna besides the course environment (pip install optuna), like the lab itself.
import sys
from io import StringIO

import matplotlib.pyplot as plt
import numpy as np
import optuna
from scipy.stats import loguniform, rv_discrete
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 (GridSearchCV, HalvingRandomSearchCV, RandomizedSearchCV,
                                     StratifiedKFold, cross_val_score, train_test_split)

# The lab's four searches on phoneme (27 evaluations of 3-fold CV each, or the same budget
# in rows for halving), drawn as best-so-far curves. Thin lines: random search and TPE
# rerun with seeds 1-4, to show how much of the difference is the luck of the seed.
optuna.logging.set_verbosity(optuna.logging.WARNING)
phoneme = fetch_openml(data_id=1489, as_frame=True)
X, y = phoneme.data, (phoneme.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)
cv = StratifiedKFold(n_splits=3, shuffle=True, random_state=42)
BUDGET = 27


def model(**params):
    return HistGradientBoostingClassifier(max_iter=500, early_stopping=True, n_iter_no_change=10,
                                          random_state=42, **params)


def log_int(low, high):
    k = np.arange(low, high + 1)
    return rv_discrete(values=(k, (1 / k) / (1 / k).sum()))


space = {'learning_rate': loguniform(0.01, 0.5), 'max_leaf_nodes': log_int(4, 128),
         'min_samples_leaf': log_int(2, 100), 'l2_regularization': loguniform(1e-3, 10),
         'max_features': [0.5, 0.75, 1.0]}
grid = {'learning_rate': [0.03, 0.1, 0.3], 'max_leaf_nodes': [8, 32, 128], 'min_samples_leaf': [5, 20, 80]}


def objective(trial):
    params = {'learning_rate': trial.suggest_float('learning_rate', 0.01, 0.5, log=True),
              'max_leaf_nodes': trial.suggest_int('max_leaf_nodes', 4, 128, log=True),
              'min_samples_leaf': trial.suggest_int('min_samples_leaf', 2, 100, log=True),
              'l2_regularization': trial.suggest_float('l2_regularization', 1e-3, 10, log=True),
              'max_features': trial.suggest_categorical('max_features', [0.5, 0.75, 1.0])}
    return cross_val_score(model(**params), X_dev, y_dev, cv=cv, scoring='roc_auc').mean()


def random_curve(seed):
    r = RandomizedSearchCV(model(), space, n_iter=BUDGET, cv=cv, scoring='roc_auc',
                           random_state=seed).fit(X_dev, y_dev)
    return np.maximum.accumulate(r.cv_results_['mean_test_score'])


def tpe_curve(seed):
    study = optuna.create_study(direction='maximize', sampler=optuna.samplers.TPESampler(seed=seed))
    study.optimize(objective, n_trials=BUDGET)
    return np.maximum.accumulate([t.value for t in study.trials])


g = GridSearchCV(model(), grid, cv=cv, scoring='roc_auc').fit(X_dev, y_dev)
curves = {'grid': np.maximum.accumulate(g.cv_results_['mean_test_score']),
          'random': random_curve(0), 'TPE': tpe_curve(0)}
halving = HalvingRandomSearchCV(model(), space, n_candidates=180, factor=3, min_resources=160,
                                cv=cv, scoring='roc_auc', random_state=1).fit(X_dev, y_dev)
default = cross_val_score(model(), X_dev, y_dev, cv=cv, scoring='roc_auc').mean()

n = np.arange(1, BUDGET + 1)
colors = {'grid': '#f0883e', 'random': '#58a6ff', 'TPE': '#3fb950'}
fig, ax = plt.subplots(figsize=(8, 3.9))
for seed in range(1, 5):
    r, t = random_curve(seed), tpe_curve(seed)
    print(f"seed {seed}: random {r[-1]:.4f} · TPE {t[-1]:.4f}", file=sys.stderr)
    ax.step(n, r, where='post', color=colors['random'], lw=0.9, alpha=0.35)
    ax.step(n, t, where='post', color=colors['TPE'], lw=0.9, alpha=0.35)
for name, c in curves.items():
    ax.step(n, c, where='post', color=colors[name], lw=2, label=f'{name} (lab run): {c[-1]:.4f}')
ax.axhline(halving.best_score_, color='gray', ls='--', lw=1.2,
           label=f'halving, same budget in rows: {halving.best_score_:.4f}')
ax.axhline(default, color='gray', ls=':', lw=1.2, label=f'defaults: {default:.4f}')
ax.set_xlim(1, BUDGET)
ax.set_ylim(0.930, 0.955)
ax.set_xlabel('evaluations (each one a 3-fold cross-validation)', fontsize=8.5)
ax.set_ylabel('best CV ROC AUC so far', fontsize=8.5)
ax.set_title('phoneme: four searches, 27 evaluations each (thin lines: random and TPE with seeds 1–4)',
             fontsize=9.5)
ax.legend(fontsize=7.5, loc='lower right')
ax.grid(alpha=0.25)
ax.tick_params(labelsize=7.5)
plt.tight_layout()

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

plt.close()