# 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.datasets import load_wine
from sklearn.metrics import accuracy_score, log_loss
from sklearn.model_selection import GridSearchCV, StratifiedKFold, train_test_split
from sklearn.neighbors import KNeighborsClassifier
from sklearn.pipeline import Pipeline
from sklearn.preprocessing import StandardScaler

X, y = load_wine(return_X_y=True)
X_train, X_test, y_train, y_test = train_test_split(
    X, y, stratify=y, random_state=15)
search = GridSearchCV(
    Pipeline([("scale", StandardScaler()), ("knn", KNeighborsClassifier())]),
    {"knn__n_neighbors": [3, 7, 11], "knn__weights": ["uniform", "distance"],
     "knn__p": [1, 2]}, scoring="accuracy",
    cv=StratifiedKFold(3, shuffle=True, random_state=15), n_jobs=1)
search.fit(X_train, y_train)
pred = search.predict(X_test)
prob = search.predict_proba(X_test)
assert pred.shape == y_test.shape
assert np.allclose(prob.sum(axis=1), 1)
assert accuracy_score(y_test, pred) > 0.7
print(search.best_params_)
print("Accuracy:", accuracy_score(y_test, pred), "log loss:", log_loss(y_test, prob))
fitted = search.best_estimator_
scaled = fitted.named_steps["scale"].transform(X_test[:2])
distance, index = fitted.named_steps["knn"].kneighbors(scaled)
print("Neighbor distances:", distance)
print("Neighbor training labels:", y_train[index])