# Source: content/notes/ml/linear-models.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 load_breast_cancer
from sklearn.linear_model import LogisticRegression
from sklearn.metrics import (balanced_accuracy_score, brier_score_loss,
                             confusion_matrix, log_loss, roc_auc_score)
from sklearn.model_selection import GridSearchCV, StratifiedKFold, train_test_split
from sklearn.pipeline import Pipeline
from sklearn.preprocessing import StandardScaler

X, y = load_breast_cancer(return_X_y=True)
X_train, X_test, y_train, y_test = train_test_split(
    X, y, test_size=0.25, stratify=y, random_state=19)
search = GridSearchCV(
    Pipeline([("scale", StandardScaler()),
              ("model", LogisticRegression(solver="lbfgs", max_iter=1500))]),
    {"model__C": [0.01, 0.1, 1.0]}, scoring="neg_log_loss",
    cv=StratifiedKFold(3, shuffle=True, random_state=19), n_jobs=1)
search.fit(X_train, y_train)
p = search.predict_proba(X_test)[:, 1]
pred = search.predict(X_test)
assert np.isfinite(p).all() and ((p >= 0) & (p <= 1)).all()
assert np.array_equal(pred, (p >= 0.5).astype(int))
assert roc_auc_score(y_test, p) > 0.85
print("Chosen C:", search.best_params_)
print({"log_loss": log_loss(y_test, p), "Brier": brier_score_loss(y_test, p),
       "ROC AUC": roc_auc_score(y_test, p),
       "balanced_accuracy": balanced_accuracy_score(y_test, pred)})
print("Confusion matrix:\n", confusion_matrix(y_test, pred))
