# Path: pysaliency/datasets.py
# def create_subset(stimuli, fixations, stimuli_indices):
#     """Create subset of stimuli and fixations using only stimuli
#     with given indices.
#     """
#     new_stimuli = stimuli[stimuli_indices]
#     if isinstance(fixations, FixationTrains):
#         fix_inds = np.in1d(fixations.train_ns, stimuli_indices)
#         new_fixations = fixations.filter_fixation_trains(fix_inds)
# 
#         index_list = list(stimuli_indices)
#         new_pos = {i: index_list.index(i) for i in index_list}
# 
#         new_fixation_train_ns = [new_pos[i] for i in new_fixations.train_ns]
#         new_fixations.train_ns = np.array(new_fixation_train_ns)
#         new_fixation_ns = [new_pos[i] for i in new_fixations.n]
#         new_fixations.n = np.array(new_fixation_ns)
#     else:
#         fix_inds = np.in1d(fixations.n, stimuli_indices)
#         new_fixations = fixations[fix_inds]
# 
#         index_list = list(stimuli_indices)
#         new_pos = {i: index_list.index(i) for i in index_list}
#         new_fixation_ns = [new_pos[i] for i in new_fixations.n]
#         new_fixations.n = np.array(new_fixation_ns)
# 
#     return new_stimuli, new_fixations


import numpy as np
from boltons.iterutils import chunked
from .datasets import create_subset
    from sklearn.model_selection import StratifiedKFold

    if split == 'train':
        folds = train_folds
    elif split == 'val':
        folds = val_folds
    elif split == 'test':
        folds = test_folds
    else:
        raise ValueError(split)

    return _get_crossval_split(stimuli, fixations, crossval_folds, included_splits=folds, random=random, stratified_attributes=stratified_attributes)


def _get_crossval_split(stimuli, fixations, split_count, included_splits, random=True, stratified_attributes=None):
    if stratified_attributes is not None:
        return _get_stratified_crossval_split(stimuli, fixations, split_count, included_splits, random=random, stratified_attributes=stratified_attributes)

    inds = list(range(len(stimuli)))
    if random:
        print("Using random shuffles for crossvalidation")
        rst = np.random.RandomState(seed=42)
        rst.shuffle(inds)
        inds = list(inds)
    size = int(np.ceil(len(inds) / split_count))
    chunks = chunked(inds, size=size)

    inds = []
    for split_nr in included_splits:
        inds.extend(chunks[split_nr])


# Path: pysaliency/datasets.py
# def create_subset(stimuli, fixations, stimuli_indices):
#     """Create subset of stimuli and fixations using only stimuli
#     with given indices.
#     """
#     new_stimuli = stimuli[stimuli_indices]
#     if isinstance(fixations, FixationTrains):
#         fix_inds = np.in1d(fixations.train_ns, stimuli_indices)
#         new_fixations = fixations.filter_fixation_trains(fix_inds)
# 
#         index_list = list(stimuli_indices)
#         new_pos = {i: index_list.index(i) for i in index_list}
# 
#         new_fixation_train_ns = [new_pos[i] for i in new_fixations.train_ns]
#         new_fixations.train_ns = np.array(new_fixation_train_ns)
#         new_fixation_ns = [new_pos[i] for i in new_fixations.n]
#         new_fixations.n = np.array(new_fixation_ns)
#     else:
#         fix_inds = np.in1d(fixations.n, stimuli_indices)
#         new_fixations = fixations[fix_inds]
# 
#         index_list = list(stimuli_indices)
#         new_pos = {i: index_list.index(i) for i in index_list}
#         new_fixation_ns = [new_pos[i] for i in new_fixations.n]
#         new_fixations.n = np.array(new_fixation_ns)
# 
#     return new_stimuli, new_fixations


# Path: pysaliency/filter_datasets.py
import numpy as np
from boltons.iterutils import chunked
from .datasets import create_subset
    from sklearn.model_selection import StratifiedKFold


    if split == 'train':
        folds = train_folds
    elif split == 'val':
        folds = val_folds
    elif split == 'test':
        folds = test_folds
    else:
        raise ValueError(split)

    return _get_crossval_split(stimuli, fixations, crossval_folds, included_splits=folds, random=random, stratified_attributes=stratified_attributes)


def _get_crossval_split(stimuli, fixations, split_count, included_splits, random=True, stratified_attributes=None):
    if stratified_attributes is not None:
        return _get_stratified_crossval_split(stimuli, fixations, split_count, included_splits, random=random, stratified_attributes=stratified_attributes)

    inds = list(range(len(stimuli)))
    if random:
        print("Using random shuffles for crossvalidation")
        rst = np.random.RandomState(seed=42)
        rst.shuffle(inds)
        inds = list(inds)
    size = int(np.ceil(len(inds) / split_count))
    chunks = chunked(inds, size=size)

    inds = []
    for split_nr in included_splits:
        inds.extend(chunks[split_nr])


    stimuli, fixations = create_subset(stimuli, fixations, inds)