# 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.cross_decomposition import PLSRegression
from sklearn.datasets import make_regression
from sklearn.linear_model import BayesianRidge, ElasticNet, Lasso, Ridge
from sklearn.metrics import mean_squared_error
from sklearn.model_selection import GridSearchCV, KFold, train_test_split
from sklearn.pipeline import Pipeline
from sklearn.preprocessing import StandardScaler

X, y = make_regression(n_samples=180, n_features=20, n_informative=5,
                       noise=18, random_state=7)
X[:, 10:15] = X[:, :5] + 0.03 * np.random.default_rng(7).normal(size=(180, 5))
X_train, X_test, y_train, y_test = train_test_split(X, y, random_state=7)
cv = KFold(n_splits=3, shuffle=True, random_state=7)
for name, estimator in [("ridge", Ridge()), ("lasso", Lasso(max_iter=15000)),
                        ("elastic", ElasticNet(l1_ratio=0.5, max_iter=15000))]:
    search = GridSearchCV(
        Pipeline([("scale", StandardScaler()), ("model", estimator)]),
        {"model__alpha": [0.1, 1.0, 10.0]}, cv=cv,
        scoring="neg_mean_squared_error", n_jobs=1)
    search.fit(X_train, y_train)
    pred = search.predict(X_test)
    assert np.isfinite(pred).all()
    print(name, search.best_params_, np.sqrt(mean_squared_error(y_test, pred)))
for name, estimator in [("Bayesian ridge", BayesianRidge()),
                        ("PLS", PLSRegression(n_components=5))]:
    pipe = Pipeline([("scale", StandardScaler()), ("model", estimator)])
    pipe.fit(X_train, y_train)
    pred = np.asarray(pipe.predict(X_test)).reshape(-1)
    assert pred.shape == y_test.shape
    print(name, np.sqrt(mean_squared_error(y_test, pred)))