
from krgee import Logistic, k_fold_cv_random_vcv

import numpy as np
import matplotlib.pyplot as plt


def plot_dataset(X, y, axes):
    plt.plot(X[:, 0][y==0], X[:, 1][y==0], "bs")
    plt.plot(X[:, 0][y==1], X[:, 1][y==1], "g^")
    plt.axis(axes)
    plt.grid(True, which='major', linestyle='--', linewidth='0.5')
    plt.xlabel(r"$x_1$", fontsize=20)
    plt.ylabel(r"$x_2$", fontsize=20, rotation=0)

def plot_predictions(clf, axes):
    x0s = np.linspace(axes[0], axes[1], 100)
    x1s = np.linspace(axes[2], axes[3], 100)
    x0, x1 = np.meshgrid(x0s, x1s)
    X = np.c_[x0.ravel(), x1.ravel()]
    y_pred = clf.predict(X).reshape(x0.shape)
    plt.contourf(x0, x1, y_pred, cmap=plt.cm.bwr , alpha=0.45)


n_updts = 20
lam = .0001
gam = 2.1
tol = 1e-4

import time
from sklearn.datasets import make_moons

X, y = make_moons(n_samples=550,
                noise=0.18,
                random_state=12038,
                shuffle=True)

num_train = int(0.65*X.shape[0])

X_train, y_train = X[:num_train], y[:num_train]
X_test, y_test = X[num_train:], y[num_train:]

vcv_train = np.eye(X_train.shape[0])
vcv_test = np.eye(X_test.shape[0])

lkrr = Logistic(kernel='rbf',
                max_iter=37, lam=lam, 
                verbose=False, tol=tol)

params = {
    'lambda' : np.logspace(-5, 3, 200, base=2),
    'gamma' : np.logspace(-5, 3, 200,  base=2),
}

best_params = k_fold_cv_random_vcv(X_train, y_train, vcv_train, 
                                   lkrr, params,
                                   folds=4, sample=200, 
                                   verbose=True, seed=12038)
print('Best params: ', best_params)
lkrr.set_params(**best_params)



# start = time.time()
lkrr.fit(X_test, y_test, vcv_test)
# print('Time taken: ', time.time() - start)
print('Log. krr loss: ', lkrr.score(X_test, y_test))

from sklearn.linear_model import LogisticRegression
log_reg = LogisticRegression()
log_reg.fit(X_train, y_train)
log_reg_loss = 1 -  log_reg.score(X_test, y_test)
print( 'Log. reg. loss: ', round(log_reg_loss, 2) )


tmp_axes =  [-1.5, 2.5, -1, 1.5]
fig, axes = plt.subplots(1, 2, figsize=(10, 4))
plt.sca(axes[0])
plot_predictions(log_reg, tmp_axes)
plot_dataset(X_test, y_test, tmp_axes)

plt.sca(axes[1])
plot_predictions(lkrr, tmp_axes)
plot_dataset(X_test, y_test, tmp_axes)
plt.tight_layout()
# plt.savefig('classification_comparisons.png', dpi=300)
