import types

from tidyfit.sklearn_extra import sklearn_stratified_split, sklearn_stratified_cv_scores


class DummyModel:
    def fit(self, X, y):
        self.major = 1 if sum(y) >= len(y) / 2 else 0

    def predict(self, X):
        return [self.major for _ in range(len(X))]


def test_sklearn_extra_with_fake_sklearn(monkeypatch):
    fake_model_selection = types.ModuleType("sklearn.model_selection")
    fake_metrics = types.ModuleType("sklearn.metrics")

    def fake_tts(X, y, test_size=0.2, random_state=None, stratify=None):
        n = len(X)
        k = int(n * (1 - test_size))
        return X[:k], X[k:], y[:k], y[k:]

    class FakeSKF:
        def __init__(self, n_splits=2, shuffle=True, random_state=None):
            self.n_splits = n_splits

        def split(self, X, y):
            n = len(X)
            half = n // 2
            yield list(range(half, n)), list(range(0, half))
            yield list(range(0, half)), list(range(half, n))

    def fake_acc(y_true, y_pred):
        return sum(int(a == b) for a, b in zip(y_true, y_pred)) / len(y_true)

    def fake_f1(y_true, y_pred, average="binary"):
        return fake_acc(y_true, y_pred)

    fake_model_selection.train_test_split = fake_tts
    fake_model_selection.StratifiedKFold = FakeSKF
    fake_metrics.accuracy_score = fake_acc
    fake_metrics.f1_score = fake_f1

    monkeypatch.setitem(
        __import__("sys").modules, "sklearn.model_selection", fake_model_selection
    )
    monkeypatch.setitem(__import__("sys").modules, "sklearn.metrics", fake_metrics)

    X = [[1], [2], [3], [4]]
    y = [0, 0, 1, 1]

    xtr, xte, ytr, yte = sklearn_stratified_split(X, y, test_size=0.5, random_state=1)
    assert len(xtr) == 2 and len(xte) == 2

    rows = sklearn_stratified_cv_scores(DummyModel(), X, y, n_splits=2, random_state=1)
    assert len(rows) == 2
    assert "accuracy" in rows[0] and "f1" in rows[0]
