from __future__ import annotations

from typing import Any
import numpy as np


def _require_sklearn() -> Any:
    try:
        from sklearn.model_selection import StratifiedKFold, train_test_split
        from sklearn.metrics import accuracy_score, f1_score
    except Exception as e:  # pragma: no cover
        raise ImportError(
            "scikit-learn is required for tidyfit.sklearn_extra. "
            "Install with: pip install 'tidyfit[sklearn]'"
        ) from e
    return StratifiedKFold, train_test_split, accuracy_score, f1_score


def sklearn_stratified_split(
    X,
    y,
    test_size: float = 0.2,
    random_state: int | None = None,
):
    """Use sklearn train_test_split with stratify=y."""
    _, train_test_split, _, _ = _require_sklearn()
    return train_test_split(
        X, y, test_size=test_size, random_state=random_state, stratify=y
    )


def sklearn_stratified_cv_scores(
    model,
    X,
    y,
    n_splits: int = 5,
    random_state: int | None = None,
):
    """Run StratifiedKFold and return accuracy/f1 per fold."""
    StratifiedKFold, _, accuracy_score, f1_score = _require_sklearn()

    X_arr = np.asarray(X)
    y_arr = np.asarray(y)
    cv = StratifiedKFold(n_splits=n_splits, shuffle=True, random_state=random_state)
    rows: list[dict[str, float]] = []

    for i, (train_idx, val_idx) in enumerate(cv.split(X_arr, y_arr), start=1):
        X_train, X_val = X_arr[train_idx], X_arr[val_idx]
        y_train, y_val = y_arr[train_idx], y_arr[val_idx]
        model.fit(X_train, y_train)
        pred = model.predict(X_val)
        rows.append(
            {
                "fold": float(i),
                "accuracy": float(accuracy_score(y_val, pred)),
                "f1": float(f1_score(y_val, pred, average="binary")),
            }
        )
    return rows
