# Path: cogspaces/utils.py
# def zip_data(data, target):
#     return {study: (data[study], target[study]) for study in data}
# 
# def unzip_data(data):
#     return {study: data[study][0] for study in data}, \
#            {study: data[study][1] for study in data}


import numpy as np
from sklearn.model_selection import GroupShuffleSplit
from cogspaces.utils import zip_data, unzip_data

    test_data : Dict[str, np.ndarray]

    train_target : Dict[str, pd.DataFrame]

    test_target : Dict[str, pd.DataFrame]
    """
    data = zip_data(data, target)
    datasets = {'train': {}, 'test': {}}

    if isinstance(test_size, (float, int)):
        test_size = {study: test_size for study in data}
    if isinstance(train_size, (float, int)):
        train_size = {study: train_size for study in data}

    splits = {}
    for study, (this_data, this_target) in data.items():
        if train_size[study] == 1.:
            train = np.arange(len(this_data))
            test = [0]
        else:
            cv = GroupShuffleSplit(n_splits=1, test_size=test_size[study],
                                   train_size=train_size[study],
                                   random_state=random_state)
            train, test = next(cv.split(X=this_data,
                                        groups=this_target['subject']))
        for fold, indices in [('train', train), ('test', test)]:
            datasets[fold][study] = this_data[indices], \
                                    this_target.iloc[indices]
    train_data, train_target = unzip_data(datasets['train'])

# Path: cogspaces/utils.py
# def zip_data(data, target):
#     return {study: (data[study], target[study]) for study in data}
# 
# def unzip_data(data):
#     return {study: data[study][0] for study in data}, \
#            {study: data[study][1] for study in data}


# Path: cogspaces/model_selection.py
import numpy as np
from sklearn.model_selection import GroupShuffleSplit
from cogspaces.utils import zip_data, unzip_data


    test_data : Dict[str, np.ndarray]

    train_target : Dict[str, pd.DataFrame]

    test_target : Dict[str, pd.DataFrame]
    """
    data = zip_data(data, target)
    datasets = {'train': {}, 'test': {}}

    if isinstance(test_size, (float, int)):
        test_size = {study: test_size for study in data}
    if isinstance(train_size, (float, int)):
        train_size = {study: train_size for study in data}

    splits = {}
    for study, (this_data, this_target) in data.items():
        if train_size[study] == 1.:
            train = np.arange(len(this_data))
            test = [0]
        else:
            cv = GroupShuffleSplit(n_splits=1, test_size=test_size[study],
                                   train_size=train_size[study],
                                   random_state=random_state)
            train, test = next(cv.split(X=this_data,
                                        groups=this_target['subject']))
        for fold, indices in [('train', train), ('test', test)]:
            datasets[fold][study] = this_data[indices], \
                                    this_target.iloc[indices]
    train_data, train_target = unzip_data(datasets['train'])

    test_data, test_target = unzip_data(datasets['test'])