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

# %% ../nbs/01_text_augmentation.ipynb 4
from __future__ import annotations
from functools import partial, wraps
import unidecode
import numpy as np
from tqdm import tqdm
from .utils import val2iterable
import pandas as pd

# %% auto 0
__all__ = ['remove_vnmese_accent', 'sampling_with_condition']

# %% ../nbs/01_text_augmentation.ipynb 5
def _remove_kwargs(func):
    @wraps(func)
    def wrapper(*args, **kwargs):
        kwargs.pop("apply_to_all", True)
        return func(*args, **kwargs)
    return wrapper

def _sampling_content(content,frac=1,seed=42,others=None):
    replace=frac>1
    rng = np.random.default_rng(seed)
    _len = len(content)
    idxs = rng.choice(list(range(_len)),int(frac*_len),replace=replace)
    content = content[idxs]
    if others is not None:
        others = others.iloc[idxs]
    return content,others

# %% ../nbs/01_text_augmentation.ipynb 6
@_remove_kwargs
def remove_vnmese_accent(content:np.ndarray|list, # A list or Numpy array of string
                         frac=1, # Fraction of the content to perform augmentation
                         seed=42, # Random seed
                         others=None # Metadata associating with the content
                        ):
    "Perform Vietnamese accent removal"
    content = val2iterable(content,t='nparray')
    if isinstance(content,list):
        content = np.array(content)
    content,others = _sampling_content(content,frac=frac,seed=seed,others=others)
    content = np.array([unidecode.unidecode(c) for c in tqdm(content)])
    if others is None:
        return content
    return content,others

# %% ../nbs/01_text_augmentation.ipynb 12
@_remove_kwargs
def sampling_with_condition(content:np.ndarray|list, # Numpy array of string
                              query:str, # Pandas query string for query method
                              frac=1, # Fraction of the content to perform augmentation
                              seed=42, # Random seed
                              others:pd.DataFrame=None, # Metadata (as dataframe) that you can query on
                           ):
    """
    Can perform oversampling/undersampling based on dataframe query
    
    For more information about dataframe query: https://pandas.pydata.org/docs/user_guide/indexing.html#indexing-query
    """
    if isinstance(content,list):
        content = np.array(content)
    idx_to_sample = others.query(query).index.values
    others_to_sample = others.loc[idx_to_sample].copy()
    content_to_sample=content[idx_to_sample].copy()
    return _sampling_content(content_to_sample,frac,seed,others=others_to_sample)
