# Source: content/notes/ml/imbalanced-data-and-pitfalls.md
# Independent CPU example; use the curriculum environment.
# See /notes/ml/#example-environment or /notes/deep-learning/#example-environment.

import numpy as np
from sklearn.metrics import balanced_accuracy_score, recall_score
from sklearn.model_selection import train_test_split
from sklearn.neighbors import KNeighborsClassifier

rng = np.random.default_rng(132)
X = rng.normal(size=(400, 6))
y = np.zeros(400, dtype=int)
y[rng.choice(400, size=40, replace=False)] = 1
source_ids = np.arange(len(y))
duplicated_ids = np.r_[source_ids, np.repeat(source_ids[y == 1], 9)]
train_rows, test_rows = train_test_split(np.arange(len(duplicated_ids)), test_size=0.3,
                                        stratify=y[duplicated_ids], random_state=132)
leaky_train, leaky_test = duplicated_ids[train_rows], duplicated_ids[test_rows]
overlap = np.intersect1d(leaky_train, leaky_test)
assert len(overlap) > 0
leaky = KNeighborsClassifier(n_neighbors=1).fit(X[leaky_train], y[leaky_train])
leaky_prediction = leaky.predict(X[leaky_test])

honest_train, honest_test = train_test_split(source_ids, test_size=0.3, stratify=y, random_state=133)
resampled_train = np.r_[honest_train, np.repeat(honest_train[y[honest_train] == 1], 9)]
assert set(resampled_train).isdisjoint(honest_test)
honest = KNeighborsClassifier(n_neighbors=1).fit(X[resampled_train], y[resampled_train])
honest_prediction = honest.predict(X[honest_test])
print("Leaked source observations:", len(overlap))
print("Leaky balanced accuracy/recall:", balanced_accuracy_score(y[leaky_test], leaky_prediction),
      recall_score(y[leaky_test], leaky_prediction))
print("Honest balanced accuracy/recall:", balanced_accuracy_score(y[honest_test], honest_prediction),
      recall_score(y[honest_test], honest_prediction))
assert recall_score(y[leaky_test], leaky_prediction) > 0.95