# AUTOGENERATED! DO NOT EDIT! File to edit: ../../nbs/01u_utils.splits.ipynb.

# %% auto 0
__all__ = ['random', 'from_fn']

# %% ../../nbs/01u_utils.splits.ipynb 2
from fastcore.all import *
import random as _random
from collections import defaultdict
from .misc import random_local_seed

# %% ../../nbs/01u_utils.splits.ipynb 3
def random(items, probs, seed=None):
    # Calculate split indexes
    n = len(items)
    p = [int(round(prob * n)) for prob in probs]  # convert percentage to absolute and round
    p[p.index(max(p))] += n - sum(p)  # adjusts the largest split to ensure the total sum matches the length of items
    # Shuffle items with the given seed
    with random_local_seed(seed):
        shuffled = items[:]
        _random.shuffle(shuffled)
    # Create cumulative split indexes and perform the split
    split_points = [sum(p[:i+1]) for i in range(len(p))]
    splits = [shuffled[start:end] for start, end in zip([0] + split_points[:-1], split_points)]
    return splits

# %% ../../nbs/01u_utils.splits.ipynb 5
def from_fn(items, fn):
    "`fn` should return the index for each subset"
    splits = defaultdict(list)
    for item in items: splits[fn(item)].append(item)
    return list(zip(*sorted(splits.items())))[1]
