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

plt.rcParams.update({'text.color': 'gray', 'axes.labelcolor': 'gray', 'axes.edgecolor': 'gray',
                     'xtick.color': 'gray', 'ytick.color': 'gray'})

# The lab's comparison, drawn: recall of each candidate on each of the SAME five folds
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)
cv = StratifiedKFold(n_splits=5, shuffle=True, random_state=42)

model = make_pipeline(StandardScaler(), LogisticRegression(max_iter=1000))
knn = make_pipeline(StandardScaler(), KNeighborsClassifier())
search_lr = GridSearchCV(model, {'logisticregression__C': np.logspace(-3, 3, 13),
                                 'logisticregression__class_weight': [None, 'balanced']},
                         cv=cv, scoring=['recall', 'roc_auc'], refit='recall').fit(X_train, y_train)
search_knn = GridSearchCV(knn, {'kneighborsclassifier__n_neighbors': [1, 3, 5, 7, 9, 11, 15, 21, 31],
                                'kneighborsclassifier__weights': ['uniform', 'distance']},
                          cv=cv, scoring=['recall', 'roc_auc'], refit='recall').fit(X_train, y_train)
candidates = {'logistic\ndefault': model, 'logistic\ntuned': search_lr.best_estimator_,
              'k-NN\ndefault': knn, 'k-NN\ntuned': search_knn.best_estimator_}
recall = {n: cross_validate(m, X_train, y_train, cv=cv, scoring='recall')['test_score']
          for n, m in candidates.items()}

fig, (left, right) = plt.subplots(1, 2, figsize=(9, 3.5), gridspec_kw={'width_ratios': [1.5, 1]})
names = list(recall)
colors = ['#58a6ff', '#f0883e', '#3fb950', '#ff7b72', '#bc8cff']
for f in range(5):
    left.plot(range(4), [recall[n][f] for n in names], 'o-', color=colors[f], ms=4, lw=1.2,
              alpha=0.85, label=f'fold {f}')
for i, n in enumerate(names):
    left.text(i, 0.855, f'{recall[n].mean():.3f}\n± {recall[n].std():.3f}', ha='center', fontsize=7.5)
left.set_xticks(range(4), names, fontsize=8)
left.set_ylim(0.84, 1.015)
left.set_xlim(-0.4, 3.4)
left.set_ylabel('recall on the validation fold', fontsize=8.5)
left.legend(fontsize=7, ncol=5, loc='upper center', bbox_to_anchor=(0.5, -0.17), frameon=False)
left.set_title('Same five folds, four candidates (one line per fold)', fontsize=9)
left.tick_params(labelsize=8)

d = recall['logistic\ntuned'] - recall['k-NN\ntuned']
tumours = np.round(d * 34).astype(int)          # 34 malignant tumours in each validation fold
right.bar(range(5), tumours, color=['#58a6ff' if t > 0 else '#f0883e' for t in tumours], width=0.6)
right.axhline(0, color='gray', lw=0.8)
right.axhspan(tumours.mean() - tumours.std(), tumours.mean() + tumours.std(), color='gray', alpha=0.15,
              lw=0)
right.axhline(tumours.mean(), color='gray', lw=1, ls='--')
right.text(-0.45, 3.8, f'mean {tumours.mean():+.1f} ± {tumours.std():.1f} tumours per fold\n'
           f'(recall {d.mean():+.3f} ± {d.std():.3f})', fontsize=7.5, ha='left', va='top')
right.set_xticks(range(5), [f'fold {f}' for f in range(5)], fontsize=8)
right.set_ylabel('tumours caught: logistic − k-NN', fontsize=8.5)
right.set_ylim(-2.5, 4)
right.set_title(f'Tuned logistic vs tuned k-NN: ahead in {(d > 0).sum()} folds,\n'
                f'behind in {(d < 0).sum()} — a tie within the noise', fontsize=9)
right.tick_params(labelsize=8)

plt.tight_layout()

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

plt.close()