from collections import namedtuple

import numpy as np
from sklearn.model_selection import StratifiedShuffleSplit

# FoldDataset = namedtuple("FoldDataset", ["x", "y"])

TRAIN_SIZE = 109


def folds(dataset="B", train_size=TRAIN_SIZE,  n_splits=3, rng=None):
    if dataset != "B":
        raise NotImplementedError()
    if rng is None:
        rng = np.random.RandomState(0)

    kf = StratifiedShuffleSplit(n_splits=n_splits,
                                train_size=0.9, random_state=rng)
    indices = np.arange(train_size)
    classes = np.array(([0] * 50) + ([1] * 50) + ([2] * 9))
    return kf.split(indices, classes)


# def k_folds(x, y, n_splits=3, train=0.4, seed=None):
#     if seed:
#         np.random.seed(seed)
#     data = np.stack([x, y], axis=-1).copy()

#     h = len(data)

#     tr = int(h * train)
#     val = tr + ((h - tr) // 2)

#     for i in range(n_splits):
#         np.random.shuffle(data)

#         train = data[:tr]
#         valid = data[tr:val]
#         test = data[val:]

#         yield [FoldDataset(a[..., 0], a[..., 1]) for a in [train, valid, test]]
