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

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

# %% ../../src/dataset/dataset_helper.ipynb 3
from ..imports import *
from ..config_loader import *

# %% ../../src/dataset/dataset_helper.ipynb 5
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)
    
    num  = comp.nonzero().squeeze().numel()  
    return bool(num)

# %% ../../src/dataset/dataset_helper.ipynb 6
def check_duplicates_in_dataset_python(xs, dataset):
    cnt = 0
     
    raise NotImplementedError("")
    
    # f = lambda x: int(check_duplicate_in_dataset(x, dataset))
    # res = async_loop_consumer(f, xs)
    # cnt = sum(res)
    
    comp = []
    
    for i,x in enumerate(xs):
        if check_duplicate_in_dataset(x, dataset):
            #print(f"[INFO] Duplicate in dataset at index={i}")
            comp.append(i)
            cnt += 1           
    # print(f"[INFO] Found {cnt}/{xs.shape[0]} duplicates in dataset of {dataset.shape[0]}.")
    
    return cnt, comp

# %% ../../src/dataset/dataset_helper.ipynb 7
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]                           
    
    # except Exception as er:
    #     print("[WARNING] check_duplicates_in_dataset:", er)
    #     print("We will use python instead.")
    #     raise NotImplementedError("")
    #     # cnt, comp = check_duplicates_in_dataset_python(xs, dataset)
        
    if return_ind: return num, comp.squeeze()  #comp is [i_xs, i_dataset] pairs
    return num

# %% ../../src/dataset/dataset_helper.ipynb 11
def shuffle_tensor_dataset(x, y=None, *z):
    '''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]         
        return x[rand_indx], y[rand_indx], *(iz[rand_indx] for iz in z)
    
    return x[rand_indx]

# %% ../../src/dataset/dataset_helper.ipynb 12
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, stable=True) #e.g. gets the index that points to zero at pos [0]
    
    cum_sum = cnt.cumsum(0)
    cum_sum = torch.cat((torch.tensor([0], device=tensor.device), cum_sum[:-1]))
        
    return tensor_unique, ind_sorted[cum_sum]

# %% ../../src/dataset/dataset_helper.ipynb 13
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 14
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):
    '''Assumes `x` is tensor and `y` is tensor or numpy.'''
    
    y_type = type(y)
    assert y_type in [np.ndarray, torch.Tensor]
    
    #------------------------------
    
    if make_unique:
        x, y, *z = uniquify_tensor_dataset(x, y, *z)
        assert x.shape[0] == y.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)      
         
    #------------------------------
    
    if y_type == np.ndarray: y_uniques_temp, y_uniques_cnt =    np.unique(y, return_counts=True, axis=0)
    else:                    y_uniques_temp, y_uniques_cnt = torch.unique(y, return_counts=True, dim=0)
    
    if y_uniques is None: y_uniques = y_uniques_temp
    if   samples is None:    
        if y_type == np.ndarray: samples =    np.min(y_uniques_cnt)   # the actual balancing count
        else:                    samples = torch.min(y_uniques_cnt)
         
    ind = list()   
    for y_unique in y_uniques:
        
        if y_type == np.ndarray:
            comp    = (y==y_unique)
            indices = np.squeeze(np.nonzero(comp))
            indices = indices if indices.ndim > 0 else indices[None]
            
        else:
            comp    = torch.all(y==y_unique, dim=1)
            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)
            
        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

# %% ../../src/dataset/dataset_helper.ipynb 16
def map_old_tensor_to_new(x):
    raise DeprecationWarning("[WARNING] There really should be no more old tensors arround .... delete them")
    print("[WARNING] There really should be no more old tensors arround .... delete them")
    
    b, gc, bits, t = x.shape
    
    x = x.reshape((b, gc//3, 3, bits, t))      # [b, g-c,  bits, t] -> [b, g, c, bits, t]   
    x = torch.argmax(x, dim=2)                 # [b, g, c, bits, t]-> [b, g, bits, t]   
    
    gate = torch.concat([torch.zeros_like(x[:,:1]), x], dim=1)  # add zeros for empty token   
    gate = torch.argmax(gate, dim=1)
    
    control_target = torch.sum(x, dim=1)
    mapped_tensor  = torch.zeros_like(control_target)
    mapped_tensor[control_target==1] = -1
    mapped_tensor[control_target==2] =  1
     
    return gate * mapped_tensor   # is now [b, space, time] with elements +-gate_number
