from collections import Counter

import numpy as np
import sklearn
from sklearn.datasets import load_breast_cancer
from sklearn.model_selection import train_test_split


dataset = load_breast_cancer()
X = dataset.data
y = dataset.target
row_ids = np.arange(X.shape[0])

X_train, X_test, y_train, y_test, train_ids, test_ids = train_test_split(
    X,
    y,
    row_ids,
    test_size=0.25,
    stratify=y,
    random_state=42,
)

_, _, _, _, _, repeat_test_ids = train_test_split(
    X,
    y,
    row_ids,
    test_size=0.25,
    stratify=y,
    random_state=42,
)


def class_counts(labels):
    counts = Counter(labels)
    return {
        str(dataset.target_names[int(label)]): int(count)
        for label, count in sorted(counts.items())
    }


print(f"scikit-learn {sklearn.__version__}")
print(f"total rows: {X.shape[0]}")
print(f"train shape: {X_train.shape}")
print(f"test shape: {X_test.shape}")
print(f"train class counts: {class_counts(y_train)}")
print(f"test class counts: {class_counts(y_test)}")
print(f"test share: {len(test_ids) / len(row_ids):.3f}")
print(f"overlap rows: {len(np.intersect1d(train_ids, test_ids))}")
print(f"repeat split matches: {np.array_equal(test_ids, repeat_test_ids)}")
print(f"first five test row ids: {test_ids[:5].tolist()}")