import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
from io import StringIO
from sklearn.datasets import make_moons
from sklearn.model_selection import GridSearchCV, StratifiedKFold, cross_val_score
from sklearn.pipeline import make_pipeline
from sklearn.preprocessing import StandardScaler
from sklearn.svm import SVC

# The lesson's grid search on noisy moons: boundaries for 9 (C, gamma) pairs,
# and the cross-validated accuracy of all 36 pairs
X, y = make_moons(n_samples=300, noise=0.3, random_state=1)
cv = StratifiedKFold(n_splits=5, shuffle=True, random_state=0)
Cs, gammas = np.logspace(-2, 3, 6), np.logspace(-2, 3, 6)
grid = GridSearchCV(make_pipeline(StandardScaler(), SVC()),
                    {'svc__C': Cs, 'svc__gamma': gammas}, cv=cv).fit(X, y)
res = pd.DataFrame(grid.cv_results_)
heat = res.pivot(index='param_svc__C', columns='param_svc__gamma', values='mean_test_score').to_numpy()

fig = plt.figure(figsize=(10.2, 5.6))
gs = fig.add_gridspec(3, 5, width_ratios=[1, 1, 1, 0.42, 1.75], wspace=0.08, hspace=0.35)
pad = 0.4
x0, x1 = X[:, 0].min() - pad, X[:, 0].max() + pad
y0, y1 = X[:, 1].min() - pad, X[:, 1].max() + pad
xx, yy = np.meshgrid(np.linspace(x0, x1, 150), np.linspace(y0, y1, 150))
show_C, show_g = [1000, 1, 0.01], [0.1, 1, 100]          # top row = large C
for i, C in enumerate(show_C):
    for j, g in enumerate(show_g):
        ax = fig.add_subplot(gs[i, j])
        model = make_pipeline(StandardScaler(), SVC(C=C, gamma=g)).fit(X, y)
        acc = heat[list(Cs).index(C), list(gammas).index(g)]
        Z = model.decision_function(np.c_[xx.ravel(), yy.ravel()]).reshape(xx.shape)
        ax.contourf(xx, yy, Z > 0, levels=[-0.5, 0.5, 1.5], colors=['#f0883e', '#58a6ff'],
                    alpha=0.18).set_rasterized(True)
        ax.contour(xx, yy, Z, levels=[0], colors='gray', linewidths=1).set_rasterized(True)
        ax.scatter(X[:, 0], X[:, 1], c=np.where(y == 1, '#58a6ff', '#f0883e'), s=3, rasterized=True)
        ax.set_xticks([]); ax.set_yticks([])
        ax.set_title(f'train {model.score(X, y):.2f} · CV {acc:.2f}', fontsize=7.5)
        if j == 0:
            ax.set_ylabel(f'C = {C:g}', fontsize=8.5)
        if i == 2:
            ax.set_xlabel(f'γ = {g:g}', fontsize=8.5)

ax = fig.add_subplot(gs[:, 4])
im = ax.imshow(heat, cmap='Blues', vmin=0.45, vmax=1.15, origin='lower', aspect='auto')
for i in range(len(Cs)):
    for j in range(len(gammas)):
        v = heat[i, j]
        ax.text(j, i, f'{v:.2f}', ha='center', va='center', fontsize=7.5,
                color='#222222')
bi, bj = list(Cs).index(grid.best_params_['svc__C']), list(gammas).index(grid.best_params_['svc__gamma'])
ax.add_patch(plt.Rectangle((bj - 0.5, bi - 0.5), 1, 1, fill=False, ec='#3fb950', lw=2.2))
ax.set_xticks(range(len(gammas)), [f'{g:g}' for g in gammas], fontsize=7.5)
ax.set_yticks(range(len(Cs)), [f'{c:g}' for c in Cs], fontsize=7.5)
ax.set_xlabel('γ  (small = broad influence, large = local)', fontsize=8.5)
ax.set_ylabel('C  (small = tolerant, large = strict)', fontsize=8.5)
ax.set_title('5-fold CV accuracy of all 36 pairs\n(green: best)', fontsize=9)
fig.suptitle('RBF SVM on noisy moons: C and γ act together', fontsize=10)

buffer = StringIO()
plt.savefig(buffer, format="svg", transparent=True, bbox_inches="tight")
print(buffer.getvalue())

plt.close()