# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/21_data.ipynb.

# %% auto 0
__all__ = ['catdists', 'datasplit']

# %% ../nbs/21_data.ipynb 6
#| export


# %% ../nbs/21_data.ipynb 8
from typing import Union

# %% ../nbs/21_data.ipynb 11
#| export

# %% ../nbs/21_data.ipynb 13
try: import numpy as np
except ImportError: ...

try: import pandas as pd
except ImportError: ...

# %% ../nbs/21_data.ipynb 15
#| export


# %% ../nbs/21_data.ipynb 17
#| export


# %% ../nbs/21_data.ipynb 19
from quac import ints, array, tensor
from nlit import DROP

# %% ../nbs/21_data.ipynb 21
from .atyp import XYArray, LabelArray
# from utrc.attr import _shape
from .smpl import permute_indicies

# %% ../nbs/21_data.ipynb 24
def catdists(data: XYArray, n: int) -> LabelArray:
    '''Categorize data by distances from the origin into n classes.'''
    dist = np.linalg.norm(data, axis=1)
    return pd.qcut(dist, q=n, labels=False, duplicates=DROP)

# %% ../nbs/21_data.ipynb 26
def datasplit(data, splits: ints, retdfs: bool = False, seed: int = 3) -> tuple[Union[array, tensor], ...]:
    els = data.shape[0]
    # els = _shape(data)[0]
    # idx = idxperm(els, *splits, seed=seed)
    idx = permute_indicies(els, *splits, seed=seed)
    if retdfs: return tuple(data.iloc[i] for i in idx)
    try: return tuple(data[i] for i in idx)
    except: return idx

# %% ../nbs/21_data.ipynb 28
#| export
