# Source: content/notes/ml/hyperparameter-tuning.md
# Independent CPU example; use the curriculum environment.
# See /notes/ml/#example-environment or /notes/deep-learning/#example-environment.

import numpy as np
from sklearn.datasets import make_classification
from sklearn.linear_model import SGDClassifier
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import StandardScaler
from sklearn.metrics import accuracy_score

X, y = make_classification(n_samples=500, n_features=12, n_informative=8,
                           class_sep=1.4, random_state=21)
train, valid, y_train, y_valid = train_test_split(
    X, y, test_size=0.3, stratify=y, random_state=2)
scaler = StandardScaler().fit(train)
train, valid = scaler.transform(train), scaler.transform(valid)
trials = [{"lr": float(lr), "epochs": 0,
           "model": SGDClassifier(loss="log_loss", learning_rate="constant",
               eta0=float(lr), alpha=1e-4, random_state=8)}
          for lr in np.logspace(-4, -0.5, 9)]
work = 0
for budget, keep in [(1, 3), (3, 1), (9, 1)]:
    for trial in trials:
        while trial["epochs"] < budget:
            trial["model"].partial_fit(train, y_train, classes=np.array([0, 1]))
            trial["epochs"] += 1
            work += 1
        trial["score"] = accuracy_score(y_valid, trial["model"].predict(valid))
    trials = sorted(trials, key=lambda t: (-t["score"], t["lr"]))[:keep]
    print("cumulative budget", budget, "survivors", len(trials), "work", work)
assert work == 21
assert trials[0]["epochs"] == 9
assert np.isfinite(trials[0]["model"].coef_).all()
assert trials[0]["score"] > 0.65
print("selected learning rate", trials[0]["lr"], "validation", trials[0]["score"])