# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/20_smpl.ipynb.

# %% auto 0
__all__ = ['idxperm', 'permute_indicies', 'groupget', 'sample']

# %% ../nbs/20_smpl.ipynb 6
import random
from functools import wraps

# %% ../nbs/20_smpl.ipynb 8
from typing import Union

# %% ../nbs/20_smpl.ipynb 11
#| export

# %% ../nbs/20_smpl.ipynb 13
try: import numpy as np
except ImportError: ...

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

# %% ../nbs/20_smpl.ipynb 15
#| export


# %% ../nbs/20_smpl.ipynb 17
#| export


# %% ../nbs/20_smpl.ipynb 19
from nlit import SAMPLES
from quac import groupkey, array, tensor, deviceq, dataframe, size, nums
from asto import asten

# %% ../nbs/20_smpl.ipynb 21
from .cons import WRAPS_ASSIGN_ANNDOCS
from .seed import seedall
from .perc import range_percents

# %% ../nbs/20_smpl.ipynb 24
def idxperm(total: int, *percents: float, seed: int = 3) -> tuple[nums, ...]:
    '''Generates a permutation of indices divided into segments based on `percents` of `total`.
    
    Parameters
    ----------
    total : int
        The total number of indices to permute.
        
    *percents : float
        The percentages to divide the total indices into.
        
    seed : int, default 3
        The seed for the random number generator.
        
    Returns
    -------
    tuple[ints, ...]
        nums of indices for each segment based on `percents`.
        
    See Also
    --------
    idxstep : Generates a step of indices based on `percents` of `total`.
    permute_indicies : Alias for `idxperm`.
    '''
    idx = list(range(total))
    seedall(seed)
    random.shuffle(idx)
    return tuple(idx[a:b] for a, b in range_percents(total, *percents))

@wraps(idxperm, assigned=WRAPS_ASSIGN_ANNDOCS)
def permute_indicies(total: int, *percents: float, seed: int = 3) -> tuple[nums, ...]:
    '''Alias for `idxperm`.'''
    return idxperm(total, *percents, seed=seed)

# %% ../nbs/20_smpl.ipynb 26
def groupget(
    df: dataframe,
    group: groupkey, 
    groupby: str = SAMPLES,
    index: str = SAMPLES, 
    set_index: bool = True
) -> array:
    '''Extracts a subgroup from a DataFrame based on a specified grouping key and 
    optionally sets a new index for the subgroup.
    
    Parameters
    ----------
    df : dataframe
        The DataFrame from which to extract the subgroup.
        
    group : groupkey
        The key of the group to extract.
        
    groupby : str, default: 'samples'
        The column name to group the DataFrame by.
        
    index : str, default: 'samples'
        The column name to set as the new index of the subgroup if `set_index` is True.
    set_index : bool, default: True
        Determines whether to set the specified column as the new index for the subgroup.
        
    Returns
    -------
    array
        The extracted subgroup as a DataFrame. If `set_index` is True, the DataFrame will 
        have its index set to the specified column.
    '''
    
    sub = df.groupby(groupby).get_group(group)
    if set_index: sub = sub.set_index(index)        
    return sub

# %% ../nbs/20_smpl.ipynb 28
def sample(
    df: dataframe, 
    
    group: groupkey, 
    groupby: str = SAMPLES,
    index: str = SAMPLES, 
    set_index: bool = True,
        
    size: size = (100, ), 
    replace: bool = False,
    
    cpu: bool = False, 
    mps: bool = False, 
    cuda: bool = False, 
    device: deviceq = None,
) -> Union[array, tensor]:
    '''Samples entries from a specific group in a DataFrame, with options for setting the index, 
    sampling size, replacement, and tensor conversion.
    
    Parameters
    ----------
    df : dataframe
        The DataFrame from which to sample.
        
    group : groupkey
        The key of the group to sample from.
        
    groupby : str, default: 'sampels'
        The column name to group the DataFrame by for sampling.
        
    index : str, default: 'sampels'
        The column name to set as the index of the sample, if `set_index` is True.
        
    set_index : bool, default: True
        Determines whether to set the specified column as the index for the sampled entries.
    
    size : size, default: (100, )
        The size of the sample to draw from the group.
        
    replace : bool, default: False
        Whether to sample with replacement.
        
    cpu : bool, default: False
        Whether to convert the sample to a tensor on the CPU.
        
    mps : bool, default: False
        Whether to convert the sample to a tensor on the MPS device.
        
    cuda : bool, default: False
        Whether to convert the sample to a tensor on the CUDA device.
        
    device : ptdevice, optional
        Specifies the device to which the tensor should be moved. 
        Overrides `cpu`, `mps`, and `cuda` if not None.
        
    Returns
    -------
    Union[array, tensor]
        The sampled entries as a DataFrame or a tensor, depending on the device parameters.
    '''
    
    if index is None: index = groupby        
    sub = groupget(df, group, groupby, index, set_index)    
    
    idx = np.arange(sub.shape[0])
    sub = sub[np.random.choice(idx, size=size, replace=replace)]
        
    if any((device, cpu, mps, cuda)):
        return asten(sub, device=device, cpu=cpu, mps=mps, cuda=cuda)    
    return sub

# %% ../nbs/20_smpl.ipynb 30
#| export
