"""Some comonly used functions for datasets."""

# AUTOGENERATED! DO NOT EDIT! File to edit: ../../src/dataset/dataset_helper.ipynb.

# %% auto #0
__all__ = ['check_duplicate_in_dataset', 'check_duplicates_in_dataset', 'shuffle_tensor_dataset', 'get_unique_elements_indices',
           'uniquify_tensor_dataset', 'balance_tensor_dataset']

# %% ../../src/dataset/dataset_helper.ipynb #06272f6f-b4e3-4504-a90a-feebbf6ad821
from ..imports import *
from ..utils.config_loader import *
from ..utils.async_fn import run_parallel_jobs

# %% ../../src/dataset/dataset_helper.ipynb #36032308-bd0e-4409-9db0-9d89fc258e5a
def check_duplicate_in_dataset(x, dataset):
    """Check if 'x' is in 'dataset'"""
    # x        ...  [   *]
    # dataset  ...  [n, *]
    
    comp = (dataset==x)  
    comp = torch.reshape(comp, [comp.shape[0], -1])  
    comp = torch.all(comp, dim=1)
    return comp.any().item()         

# %% ../../src/dataset/dataset_helper.ipynb #2a8886b0-2b43-4c06-9992-b859d11e698e
def check_duplicates_in_dataset(xs, dataset, return_ind=False, invert=False):
    '''
    Checks if `xs` is are `dataset`. Boolean `invert` changes if we count duplicates (False) or ones that are not in dataset (True). 
    Uses `torch.vmap` which copies `dataset` for every element in `xs`.
    '''
    
    # x        ...  [c, *]
    # dataset  ...  [n, *]
        
    def get_comp(x, dataset):
        comp = (dataset==x) 
        comp = torch.reshape(comp, [comp.shape[0], -1])  
        comp = torch.all(comp, dim=1)  
        return comp 
            
    get_comp_map = torch.vmap(get_comp, in_dims=(0, None)) 
    
    # try: #could run out of memory bcs dataset is copied each time
    
    comp = get_comp_map(xs, dataset)
    
    if invert: comp = torch.all(comp==False, dim=1) 
    else:      comp = torch.any(comp, dim=1)
    
    comp = comp.nonzero()
    num  = comp.shape[0]                           
    
    if return_ind: return num, comp.squeeze()  #comp is [i_xs, i_dataset] pairs
    return num

# %% ../../src/dataset/dataset_helper.ipynb #9ea47154-d7c2-435b-8d08-0999c744af1a
def shuffle_tensor_dataset(x, y=None, *z, cpu_copy=True):
    '''Assumes numpy or tensor objects with same length.'''
    rand_indx = torch.randperm(x.shape[0])
    
    if exists(y):
        assert x.shape[0] == y.shape[0]       
        for iz in z: assert x.shape[0] == iz.shape[0]        


        if cpu_copy:

            def _cpu_array_index(var):                   
                if type(var) == np.ndarray:
                    var = var[rand_indx]
                else:
                    device = var.device
                    var = var.to("cpu")
                    var = var[rand_indx]
                    var[:] = var.to(device)                   
                return var

            x = _cpu_array_index(x)
            y = _cpu_array_index(y)
            z = (_cpu_array_index(iz) for iz in z)
            
            return x, y, *z

        else:
            return x[rand_indx], y[rand_indx], *(iz[rand_indx] for iz in z)
    
    return x[rand_indx]

# %% ../../src/dataset/dataset_helper.ipynb #c8005564-1f96-41da-94f5-16d2f9f48b2a
def get_unique_elements_indices(tensor):
    '''Returns indices of unique_elements in `tensor`.'''
    tensor_unique, ptrs, cnt = torch.unique(tensor, dim=0, return_inverse=True, return_counts=True)
    _, ind_sorted = torch.sort(ptrs, dim=0, stable=True) #e.g. gets the index that points to zero at pos [0]
    
    cum_sum = cnt.cumsum(dim=0)
    cum_sum = torch.cat([torch.tensor([0], device=tensor.device), cum_sum[:-1]], dim=0)

    idx = ind_sorted[cum_sum].cpu()
    
    return tensor[idx], idx

# %% ../../src/dataset/dataset_helper.ipynb #a9295de7-54e1-435b-aeea-26746019dd95
def uniquify_tensor_dataset(x, y=None, *z):
    '''`x` has to be tensor, assumes numpy or tensor obj for `y` and `z`'''
    x, x_idx = get_unique_elements_indices(x)
    x_idx    = x_idx.cpu()
    
    if exists(y):
        assert x.shape[0] == y[x_idx].shape[0]   
        for iz in z: assert x.shape[0] == iz[x_idx].shape[0]
        return x, y[x_idx], *(iz[x_idx] for iz in z)
    
    return x

# %% ../../src/dataset/dataset_helper.ipynb #e63e9933-a2ce-47b3-a957-58705e060fc7
def balance_tensor_dataset(x, y, *z, samples: int=None, make_unique: bool=True, y_uniques=None, shuffle_lables: bool=True, add_balance_fn: callable=None, njobs=1):
    '''Assumes `x` is tensor and `y` is tensor or numpy.'''
    
    y_type = type(y)
    assert y_type in [np.ndarray, torch.Tensor]

    print(f" - balance_tensor_dataset, {njobs=}, number of samples={x.shape[0]}")
    
    #------------------------------
    
    if make_unique:
        x, y, *z = uniquify_tensor_dataset(x, y, *z)
        assert x.shape[0] == y.shape[0]

        print(f" - uniquify_tensor_dataset, number of samples now {x.shape[0]}")
        
        #bcs unique sorts, we need to shuffle the dataset before picking the first 'samples' entries
        x, y, *z = shuffle_tensor_dataset(x, y, *z)      
         
    #------------------------------

    search_y = y_uniques if exists(y_uniques) else y
    
    if y_type == np.ndarray: _, y_ptrs, y_uniques_cnt =    np.unique(search_y, return_counts=True, return_inverse=True, axis=0)
    else:                    _, y_ptrs, y_uniques_cnt = torch.unique(search_y, return_counts=True, return_inverse=True, dim=0)
    
    if not exists(samples):    
        if y_type == np.ndarray: samples =    np.min(y_uniques_cnt)   # the actual balancing count
        else:                    samples = torch.min(y_uniques_cnt)
         
    print(f" - balancing")
    # ToDo: make parallel 
    
    ind = list()   
    # for y_unique in tqdm(y_uniques, total=y_uniques.shape[0]):   
    for y_ptr_index in tqdm(range(y_uniques_cnt.shape[0]), total=y_uniques_cnt.shape[0]):
        
        if y_type == np.ndarray:
            comp    = (y_ptrs==y_ptr_index)
            indices = np.squeeze(np.nonzero(comp))
            indices = indices if indices.ndim > 0 else indices[None]
            
        else:            
            comp = (y_ptrs==y_ptr_index)
            
            indices = comp.nonzero().squeeze() #.cpu()
            indices = indices if indices.dim() > 0 else indices[None]

        #special add balncing, e.g., for circuit length
        if add_balance_fn is not None: indices = add_balance_fn(indices, x, y, *z)

        if not y_type == np.ndarray: indices = indices.cpu()

        indices = shuffle_tensor_dataset(indices) 

        #fixes bug: shuffle_tensor_dataset removes dim if numpy array only has 1 element! 
        if y_type == np.ndarray: indices = indices if indices.ndim > 0 else indices[None]
        else:                    indices = indices if indices.dim() > 0 else indices[None]
        
        ind.append(indices[:samples]) #limit samples
        
    if y_type == np.ndarray: ind = np.concatenate(ind, axis=0)
    else:                    ind = torch.cat(ind, dim=0)
    
    xb = x[ind]
    yb = y[ind]
    zb = (iz[ind] for iz in z)
        
    #shuffle again bcs we are sorted by label; normaly we don't care and we shuffle with dataset obj anyways, but needed if valid split afterwards
    if shuffle_lables: xb, yb, *zb = shuffle_tensor_dataset(xb, yb, *zb) 
        
    return xb, yb, *zb
