# Source: content/notes/ml/probabilistic-and-instance-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.dummy import DummyRegressor
from sklearn.metrics import mean_squared_error
from sklearn.model_selection import GridSearchCV, KFold, train_test_split
from sklearn.neighbors import KNeighborsRegressor
from sklearn.pipeline import make_pipeline
from sklearn.preprocessing import StandardScaler

rng = np.random.default_rng(16)
X = rng.uniform(-3, 3, size=(240, 1))
y = np.sin(X[:, 0]) + rng.normal(0, 0.12, size=240)
X_train, X_test, y_train, y_test = train_test_split(X, y, random_state=16)
search = GridSearchCV(make_pipeline(StandardScaler(), KNeighborsRegressor()),
    {"kneighborsregressor__n_neighbors": [3, 9, 17],
     "kneighborsregressor__weights": ["uniform", "distance"]},
    scoring="neg_mean_squared_error", cv=KFold(3, shuffle=True, random_state=16))
search.fit(X_train, y_train)
pred = search.predict(X_test)
baseline = DummyRegressor().fit(X_train, y_train).predict(X_test)
assert mean_squared_error(y_test, pred) < mean_squared_error(y_test, baseline)
assert np.isfinite(pred).all()
print("RMSE:", np.sqrt(mean_squared_error(y_test, pred)))
print("Outside training range:", search.predict(np.array([[-6.0], [6.0]])))