import sys
import matplotlib.pyplot as plt
import numpy as np
from io import StringIO
from sklearn.base import clone
from sklearn.datasets import fetch_openml
from sklearn.linear_model import LogisticRegression
from sklearn.metrics import average_precision_score
from sklearn.model_selection import GridSearchCV, StratifiedKFold, train_test_split
from sklearn.neighbors import KNeighborsClassifier
from sklearn.pipeline import make_pipeline
from sklearn.preprocessing import PolynomialFeatures, StandardScaler

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

# The lesson's model comparison on the mammography data: three families, each tuned by
# GridSearchCV on the same 5 folds, the family with the best best_score_ wins.
mammo = fetch_openml(data_id=310, as_frame=True)
X, y = mammo.data, (mammo.target == '1').astype(int)
X_dev, X_test, y_dev, y_test = train_test_split(X, y, test_size=0.2, stratify=y, random_state=0)
cv = StratifiedKFold(n_splits=5, shuffle=True, random_state=0)
searches = {
    'logistic': GridSearchCV(
        make_pipeline(StandardScaler(), LogisticRegression(max_iter=1000)),
        {'logisticregression__C': np.logspace(-3, 3, 7),
         'logisticregression__class_weight': [None, 'balanced']},
        cv=cv, scoring='average_precision'),
    'logistic degree 2': GridSearchCV(
        make_pipeline(PolynomialFeatures(2), StandardScaler(), LogisticRegression(max_iter=1000)),
        {'logisticregression__C': np.logspace(-3, 3, 7)},
        cv=cv, scoring='average_precision'),
    'k-NN': GridSearchCV(
        make_pipeline(StandardScaler(), KNeighborsClassifier()),
        {'kneighborsclassifier__n_neighbors': [5, 15, 25, 35, 51, 75, 101],
         'kneighborsclassifier__weights': ['uniform', 'distance']},
        cv=cv, scoring='average_precision'),
}
naive = max(s.fit(X_dev, y_dev).best_score_ for s in searches.values())

# Nested CV, repeated with 10 different outer shuffles
estimates, fold_scores, wins = [], [], {name: 0 for name in searches}
for seed in range(1, 11):
    outer = StratifiedKFold(n_splits=5, shuffle=True, random_state=seed)
    scores = []
    for tr, va in outer.split(X_dev, y_dev):
        fitted = {n: clone(s).fit(X_dev.iloc[tr], y_dev.iloc[tr]) for n, s in searches.items()}
        winner = max(fitted, key=lambda n: fitted[n].best_score_)
        wins[winner] += 1
        scores.append(average_precision_score(y_dev.iloc[va],
                                              fitted[winner].predict_proba(X_dev.iloc[va])[:, 1]))
    fold_scores.append(scores)
    estimates.append(np.mean(scores))
estimates, fold_scores = np.array(estimates), np.array(fold_scores)
print(round(naive, 4), estimates.round(4), fold_scores.min().round(3), fold_scores.max().round(3), wins,
      file=sys.stderr)                # the numbers quoted in the text

fig, ax = plt.subplots(figsize=(8, 4.2))
rows = np.arange(1, 11)
for r, fs in zip(rows, fold_scores):
    ax.scatter(fs, [r] * 5, s=10, color='gray', alpha=0.45, lw=0)
ax.scatter(estimates, rows, s=40, color='#58a6ff', zorder=3, label='nested estimate (mean of 5 outer folds)')
ax.scatter([], [], s=10, color='gray', alpha=0.45, label='one outer fold')
ax.axvline(naive, color='#ff7b72', lw=1.6,
           label=f'best_score_ of the winner, {naive:.3f}')
ax.axvline(estimates.mean(), color='#58a6ff', lw=1.2, ls='--',
           label=f'mean of the 10 nested estimates, {estimates.mean():.3f}')
ax.set_yticks(rows, [f'shuffle {r}' for r in rows], fontsize=7.5)
ax.invert_yaxis()
ax.set_xlabel('average precision on the outer validation fold', fontsize=8.5)
ax.tick_params(labelsize=8)
ax.legend(fontsize=7.5, loc='upper center', bbox_to_anchor=(0.5, -0.17), ncol=2, frameon=False)
ax.set_xlim(0.52, 0.86)
won = ', '.join(f'{n} {k}' for n, k in wins.items())
ax.set_title(f'Nested CV of the whole choice, 10 outer shuffles: {estimates.min():.3f} to '
             f'{estimates.max():.3f}; {(estimates < naive).sum()} of 10 below {naive:.3f}\nwinner in the 50 outer folds: {won}',
             fontsize=9)

plt.tight_layout()

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

plt.close()