import matplotlib.pyplot as plt
import numpy as np
from io import StringIO
from sklearn.datasets import load_breast_cancer
from sklearn.ensemble import RandomForestClassifier
from sklearn.feature_selection import SelectKBest, f_classif
from sklearn.linear_model import LogisticRegression
from sklearn.model_selection import GroupKFold, KFold, StratifiedKFold, cross_val_score
from sklearn.neighbors import KNeighborsClassifier
from sklearn.pipeline import make_pipeline
from sklearn.preprocessing import StandardScaler

cv = StratifiedKFold(n_splits=5, shuffle=True, random_state=0)
rows = []   # (label, reported, honest)

# 1. Feature selection before CV — 100 rows, 5000 columns of pure noise
rng = np.random.default_rng(5)
X, y = rng.normal(size=(100, 5000)), rng.integers(0, 2, 100)
X_sel = SelectKBest(f_classif, k=20).fit_transform(X, y)
leaked = cross_val_score(LogisticRegression(max_iter=1000), X_sel, y, cv=cv).mean()
honest = cross_val_score(make_pipeline(SelectKBest(f_classif, k=20), LogisticRegression(max_iter=1000)),
                         X, y, cv=cv).mean()
rows.append(('SelectKBest before CV\n(pure noise)', leaked, honest))

# 2. Scaling before CV — breast cancer: the same mistake, a much smaller leak
X, y = load_breast_cancer(return_X_y=True)
leaked = cross_val_score(LogisticRegression(max_iter=1000), StandardScaler().fit_transform(X), y, cv=cv).mean()
honest = cross_val_score(make_pipeline(StandardScaler(), LogisticRegression(max_iter=1000)), X, y, cv=cv).mean()
rows.append(('StandardScaler before CV\n(breast cancer)', leaked, honest))

# 3. Near-duplicate rows — pure noise, every row has a jittered twin
rng = np.random.default_rng(4)
X, y = rng.normal(size=(300, 10)), rng.integers(0, 2, 300)
X2 = np.vstack([X, X + rng.normal(0, 0.01, X.shape)])
y2, twin = np.concatenate([y, y]), np.concatenate([np.arange(300), np.arange(300)])
knn = KNeighborsClassifier(n_neighbors=1)
leaked = cross_val_score(knn, X2, y2, cv=KFold(5, shuffle=True, random_state=0)).mean()
honest = cross_val_score(knn, X2, y2, cv=GroupKFold(5), groups=twin).mean()
rows.append(('near-duplicate rows\n(pure noise)', leaked, honest))

# 4. Group leakage — 60 patients x 10 exams, each patient with a strong "fingerprint"
def patients(n_pat, seed, per=10):
    r = np.random.default_rng(seed)
    sick = r.integers(0, 2, n_pat)
    fingerprint = r.normal(0, 3, (n_pat, 10))
    pid = np.repeat(np.arange(n_pat), per)
    X = fingerprint[pid] + r.normal(0, 0.5, (n_pat * per, 10))
    X[:, 0] = 2.0 * sick[pid] + r.normal(0, 1, n_pat * per)   # the real (weaker) signal
    return X, sick[pid], pid
X, y, pid = patients(60, seed=1)
rf = RandomForestClassifier(n_estimators=200, random_state=0)
leaked = cross_val_score(rf, X, y, cv=KFold(5, shuffle=True, random_state=0)).mean()
honest = cross_val_score(rf, X, y, cv=GroupKFold(5), groups=pid).mean()
rows.append(('same patient in train and\nvalidation (random KFold)', leaked, honest))

labels = [r[0] for r in rows]
leaked = np.array([r[1] for r in rows]) * 100
honest = np.array([r[2] for r in rows]) * 100
pos = np.arange(len(rows))[::-1]

fig, ax = plt.subplots(figsize=(6.8, 3.9))
ax.barh(pos + 0.19, leaked, height=0.36, color='#ff7b72', label='reported (leaky)')
ax.barh(pos - 0.19, honest, height=0.36, color='#3fb950', label='honest')
for p, l, h in zip(pos, leaked, honest):
    ax.text(l + 0.6, p + 0.19, f'{l:.1f}%', va='center', fontsize=8)
    ax.text(h + 0.6, p - 0.19, f'{h:.1f}%', va='center', fontsize=8)
ax.set_yticks(pos, labels, fontsize=8)
ax.axvline(50, color='gray', ls='--', lw=1)
ax.set_xlim(40, 105)
ax.set_xlabel('cross-validated accuracy (%)', fontsize=9)
ax.legend(fontsize=8, loc='upper center', bbox_to_anchor=(0.5, -0.2), ncol=2, frameon=False)
ax.set_title('Four leaks, measured: same data and model, only the validation changes', fontsize=10)

plt.tight_layout()

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

plt.close()