# 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 make_regression
from sklearn.dummy import DummyRegressor
from sklearn.impute import SimpleImputer
from sklearn.linear_model import LinearRegression
from sklearn.metrics import mean_absolute_error, mean_squared_error, r2_score
from sklearn.model_selection import train_test_split
from sklearn.pipeline import make_pipeline
from sklearn.preprocessing import StandardScaler

X, y = make_regression(n_samples=450, n_features=6, n_informative=4,
                       noise=15, random_state=12)
X_train, X_test, y_train, y_test = train_test_split(
    X, y, test_size=0.25, random_state=12)
model = make_pipeline(SimpleImputer(strategy="median"), StandardScaler(),
                      LinearRegression())
model.fit(X_train, y_train)
pred = model.predict(X_test)
baseline = DummyRegressor().fit(X_train, y_train).predict(X_test)
assert pred.shape == y_test.shape and np.isfinite(pred).all()
assert mean_squared_error(y_test, pred) < mean_squared_error(y_test, baseline)
print({"MAE": mean_absolute_error(y_test, pred),
       "RMSE": np.sqrt(mean_squared_error(y_test, pred)),
       "R2": r2_score(y_test, pred)})
print("Test residual mean:", np.mean(y_test - pred))