from __future__ import annotations

import numpy as np


def kfold_indices(n_samples: int, n_splits: int = 5, random_state: int | None = None):
    """Yield (train_idx, val_idx) tuples."""
    if n_splits < 2:
        raise ValueError("n_splits must be >=2")
    rng = np.random.default_rng(random_state)
    idx = np.arange(n_samples)
    rng.shuffle(idx)
    folds = np.array_split(idx, n_splits)
    for i in range(n_splits):
        val = folds[i]
        train = np.concatenate([folds[j] for j in range(n_splits) if j != i])
        yield train, val


def stratified_kfold_indices(
    y: np.ndarray,
    n_splits: int = 5,
    random_state: int | None = None,
):
    """Yield stratified (train_idx, val_idx) tuples for classification."""
    if n_splits < 2:
        raise ValueError("n_splits must be >=2")
    y = np.asarray(y)
    rng = np.random.default_rng(random_state)

    by_class: dict[int, np.ndarray] = {}
    for cls in np.unique(y):
        idx = np.where(y == cls)[0]
        rng.shuffle(idx)
        by_class[int(cls)] = idx

    fold_bins: list[list[int]] = [[] for _ in range(n_splits)]
    for idx in by_class.values():
        chunks = np.array_split(idx, n_splits)
        for i, c in enumerate(chunks):
            fold_bins[i].extend(c.tolist())

    for i in range(n_splits):
        val = np.array(fold_bins[i], dtype=int)
        train = np.array(
            [j for k, fold in enumerate(fold_bins) if k != i for j in fold], dtype=int
        )
        yield train, val
