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

# The lab's grid search (step 6), drawn: CV accuracy for every (C, gamma) pair
data = fetch_openml(data_id=59, as_frame=True)
X, y = data.data, (data.target == 'g').astype(int)
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.25, stratify=y, random_state=42)
cv = StratifiedKFold(n_splits=5, shuffle=True, random_state=42)
Cs, gammas = [0.1, 1, 10, 100, 1000], [0.0001, 0.001, 0.01, 0.1, 1]
grid = GridSearchCV(make_pipeline(StandardScaler(), SVC()), {'svc__C': Cs, 'svc__gamma': gammas},
                    cv=cv, scoring='accuracy').fit(X_train, y_train)
res = pd.DataFrame(grid.cv_results_)
table = res.pivot(index='param_svc__C', columns='param_svc__gamma', values='mean_test_score').to_numpy()
best_std = res.loc[grid.best_index_, 'std_test_score']
tie = table >= grid.best_score_ - best_std

fig, ax = plt.subplots(figsize=(6.8, 4.4))
ax.imshow(table, cmap='Blues', vmin=0.45, vmax=1.15, origin='lower', aspect='auto')
for i in range(len(Cs)):
    for j in range(len(gammas)):
        ax.text(j, i, f'{table[i, j]:.3f}', ha='center', va='center', fontsize=8.5, color='#222222',
                fontweight='bold' if tie[i, j] else 'normal')
        if tie[i, j]:
            ax.add_patch(plt.Rectangle((j - 0.46, i - 0.46), 0.92, 0.92, fill=False, ec='#3fb950', lw=1.6))
bi, bj = Cs.index(grid.best_params_['svc__C']), gammas.index(grid.best_params_['svc__gamma'])
ax.add_patch(plt.Rectangle((bj - 0.5, bi - 0.5), 1, 1, fill=False, ec='#f0883e', lw=2.6))
ax.set_xticks(range(len(gammas)), [f'{g:g}' for g in gammas], fontsize=8)
ax.set_yticks(range(len(Cs)), [f'{c:g}' for c in Cs], fontsize=8)
ax.set_xlabel('γ', fontsize=9)
ax.set_ylabel('C', fontsize=9)
ax.set_title(f'Ionosphere, 5-fold CV accuracy (baseline 0.643)\n'
             f'orange: best, {grid.best_score_:.3f} ± {best_std:.3f} · green: within 1 std of it',
             fontsize=9.5)
plt.tight_layout()

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

plt.close()