import matplotlib.pyplot as plt
import numpy as np
from io import StringIO
from sklearn.datasets import fetch_openml
from sklearn.linear_model import LogisticRegression
from sklearn.model_selection import RepeatedStratifiedKFold, cross_val_score, 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'})

# Two pairs of candidates on the mammography data, each scored on the SAME 50 folds
# (5-fold CV repeated 10 times). One line per fold joins the two scores of a pair.
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)
rcv = RepeatedStratifiedKFold(n_splits=5, n_repeats=10, random_state=0)

def ap(est):
    return cross_val_score(est, X_dev, y_dev, cv=rcv, scoring='average_precision')

pairs = [
    ('logistic', ap(make_pipeline(StandardScaler(), LogisticRegression(max_iter=1000))),
     'logistic\nbalanced', ap(make_pipeline(StandardScaler(),
                                           LogisticRegression(max_iter=1000, class_weight='balanced')))),
    ('k-NN\nk=25, distance', ap(make_pipeline(StandardScaler(),
                                                KNeighborsClassifier(25, weights='distance'))),
     'logistic deg. 2\nC=0.1', ap(make_pipeline(PolynomialFeatures(2), StandardScaler(),
                                                   LogisticRegression(C=0.1, max_iter=1000)))),
]

fig, axes = plt.subplots(1, 3, figsize=(10, 3.9), gridspec_kw={'width_ratios': [1, 1, 1.25]})
for ax, (na, a, nb, b), title in zip(axes[:2], pairs, ['a real difference', 'a tie']):
    d = a - b
    for sa, sb, di in zip(a, b, d):
        ax.plot([0, 1], [sa, sb], color='#58a6ff' if di > 0 else '#f0883e', alpha=0.55, lw=1)
    for x, s in [(0, a), (1, b)]:
        ax.errorbar(x + (-0.12 if x == 0 else 0.12), s.mean(), yerr=s.std(), fmt='o', color='gray',
                    capsize=3, ms=4)
        ax.text(x + (-0.17 if x == 0 else 0.17), s.mean(), f'{s.mean():.3f}\n± {s.std():.3f}',
                fontsize=7.5, color='gray', va='center', ha='right' if x == 0 else 'left')
    ax.set_xticks([0, 1], [na, nb], fontsize=8)
    ax.set_xlim(-0.75, 1.75)
    ax.set_ylim(0.4, 0.88)
    ax.set_title(f'{title}: first one ahead in {(d > 0).sum()} of {len(d)} folds', fontsize=8.5)
    ax.tick_params(labelsize=8)
axes[0].set_ylabel('average precision (one line per fold)', fontsize=8)

ax = axes[2]
rng = np.random.default_rng(0)
for row, (na, a, nb, b) in zip([1, 0], pairs):
    d = a - b
    ax.scatter(d, row + rng.uniform(-0.12, 0.12, len(d)), s=12, alpha=0.6,
               color=['#58a6ff' if v > 0 else '#f0883e' for v in d])
    ax.errorbar(d.mean(), row + 0.3, xerr=d.std(), fmt='D', color='#3fb950', capsize=4, ms=5)
    ax.text(d.mean(), row + 0.42, f'difference {d.mean():+.3f} ± {d.std():.3f}', fontsize=7.5,
            color='#3fb950', ha='center')
    # what the spread of the difference would be if the two models were scored on unrelated folds
    unpaired = np.sqrt(a.std() ** 2 + b.std() ** 2)
    ax.plot([d.mean() - unpaired, d.mean() + unpaired], [row - 0.3] * 2, color='gray', lw=1.2)
    ax.text(d.mean(), row - 0.45, f'± {unpaired:.3f} if the folds were not shared', fontsize=7,
            color='gray', ha='center')
ax.axvline(0, color='gray', lw=0.8, ls='--')
ax.set_yticks([1, 0], ['logistic −\nbalanced', 'k-NN −\ndeg. 2'], fontsize=8)
ax.set_ylim(-0.75, 1.7)
ax.set_xlabel('per-fold difference in average precision', fontsize=8)
ax.set_title('subtract fold by fold: the shared difficulty cancels', fontsize=8.5)
ax.tick_params(labelsize=8)

fig.suptitle('Same 50 folds for both models of each pair (mammography, 5-fold CV × 10 repeats)',
             fontsize=10)
plt.tight_layout()

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

plt.close()