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

# The lab's grid search, drawn: recall on the training folds and on the
# validation folds for each value of C
data = load_breast_cancer(as_frame=True)
X, y = data.data, 1 - data.target                 # 1 = malignant
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, stratify=y, random_state=42)

model = make_pipeline(StandardScaler(), LogisticRegression(max_iter=1000))
cv = StratifiedKFold(n_splits=5, shuffle=True, random_state=42)
grid = GridSearchCV(model, {'logisticregression__C': np.logspace(-3, 3, 13)}, cv=cv,
                    scoring='recall', return_train_score=True).fit(X_train, y_train)
res = pd.DataFrame(grid.cv_results_)
C = res['param_logisticregression__C'].astype(float).to_numpy()
tr, va, sd = res['mean_train_score'], res['mean_test_score'], res['std_test_score']
best = grid.best_index_

fig, ax = plt.subplots(figsize=(6.8, 3.4))
ax.fill_between(C, va - sd, va + sd, color='#58a6ff', alpha=0.15, lw=0, label='validation ± 1 std across folds')
ax.plot(C, tr, 'o-', color='#f0883e', ms=4, label='training folds')
ax.plot(C, va, 'o-', color='#58a6ff', ms=4, label='validation folds')
# every C whose mean is within one std of the best is statistically a tie
tie = va >= va[best] - sd[best]
ax.axhspan(va[best] - sd[best], va[best], color='gray', alpha=0.12, lw=0)
ax.scatter(C[tie], va[tie], s=60, facecolors='none', edgecolors='#3fb950', lw=1.4, zorder=5,
           label='within 1 std of the best')
ax.annotate(f'best: C = {C[best]:g}, recall {va[best]:.3f}', xy=(C[best], va[best]),
            xytext=(0.05, 1.0), fontsize=8, arrowprops=dict(arrowstyle='->', color='gray'))
ax.text(0.0016, 0.63, 'strong penalty:\nunderfits', fontsize=8, color='gray')
ax.text(60, 1.012, 'weak penalty: training → 1.0', fontsize=7.5, color='gray', ha='center')
ax.set_xscale('log')
ax.set_xticks([1e-3, 1e-2, 1e-1, 1, 10, 100, 1000], ['0.001', '0.01', '0.1', '1', '10', '100', '1000'])
ax.set_xlabel('C  (small = strong regularization, large = weak)')
ax.set_ylabel('recall (malignant)')
ax.set_ylim(0.6, 1.04)
ax.legend(fontsize=7, loc='lower right')
ax.set_title('Grid search over C, seen from inside the folds', fontsize=10)

plt.tight_layout()

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

plt.close()