# AUTOGENERATED! DO NOT EDIT! File to edit: nbs/000b_data.validation.ipynb (unless otherwise specified).

__all__ = ['TrainValTestSplitter', 'get_splits', 'check_overlap', 'leakage_finder', 'get_predefined_splits',
           'combine_split_data']

# Cell
from ..imports import *
from ..utils import *

# Cell
from sklearn.model_selection import train_test_split

# Cell
def TrainValTestSplitter(valid_size=0.2, test_size=0.2, random_state=None, stratify=None, **kwargs):
    "Split `items` into random train, valid and test subsets using sklearn train_test_split utility."
    valid_size = valid_size / (1 - test_size)
    stratify1 = None if stratify is None else stratify
    def _inner(o, **kwargs):
        train_valid, test = train_test_split(range(len(o)), test_size=test_size, random_state=random_state, stratify=stratify1, **kwargs)
        stratify2 = None if stratify1 is None else stratify1[train_valid]
        train, valid = train_test_split(range(len(train_valid)), test_size=valid_size, random_state=random_state, stratify=stratify2, **kwargs)
        return L(L(train_valid)[train]), L(L(train_valid)[valid]), L(test)
    return _inner

def get_splits(y, valid_size=0.2, test_size=0.2, stratify=True, merge_train_valid=False, merge_train_test=False, train_perc=None,
               random_state=None, **kwargs):
    if test_size == 0: splits = L(TrainTestSplitter(test_size=valid_size, random_state=random_state, stratify=y if stratify else None)(y))
    else: splits = L(TrainValTestSplitter(valid_size=valid_size, test_size=test_size, random_state=random_state, stratify=y if stratify else None)(y))
    if merge_train_valid: splits[0] = concat(splits[0], splits[1])
    if merge_train_test: splits[0] = concat(splits[0], splits[2])
    if train_perc: splits[0] = L(np.random.choice(len(train), int(len(train) * train_perc), False))
    leakage_finder(*splits)
    return splits

def check_overlap(a, b):
    overlap = L([i for i in a if i in b])
    if overlap == []: return
    return overlap

def leakage_finder(train, val, test=None):
    if check_overlap(train, val) is not None:
        print('train-val leakage!')
        print(check_overlap(train, val), '\n')
    if test is not None:
        if check_overlap(train, test) is not None:
            print('train-test leakage!')
            print(check_overlap(train, test), '\n')
        if check_overlap(val, test) is not None:
            print('val-test leakage!')
            print(check_overlap(val, test), '\n')

# Cell
def get_predefined_splits(*xs):
    '''xs is a list with X_train, X_valid, ...'''
    splits_ = []
    start = 0
    for x in xs:
        splits_.append(L(list(np.arange(start, start + len(x)))))
        start += len(x)
    return tuple(splits_)

def combine_split_data(xs, ys=None):
    '''xs is a list with X_train, X_valid, .... ys is None or a list with y_train, y_valid, .... '''
    xs = [to3d(x) for x in xs]
    splits = get_predefined_splits(*xs)
    if ys is None: return concat(*xs), None, splits
    else: return concat(*xs), concat(*ys), splits