"""Helper functions used to balance a dataset."""

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

# %% auto #0
__all__ = ['get_tensor_gate_length', 'add_balance_fn_quantile_qc_length']

# %% ../../src/dataset/balancing.ipynb #06272f6f-b4e3-4504-a90a-feebbf6ad821
from ..imports import *
import genQC.dataset.dataset_helper as dahe

# %% ../../src/dataset/balancing.ipynb #e57ce23e-30fc-434f-9443-3cf97f507b89
def get_tensor_gate_length(clr_tensor: torch.Tensor, padding_token: int = 0) -> torch.Tensor:
    """
    Returns the gate count of a tokenized circuit.
    Make sure you use use the correct `padding_token`.
    
    """
    assert clr_tensor.dim() == 3, "[b, s, t]"
    
    red_clr_tensor = (clr_tensor != padding_token).any(dim=1) # [b, t]
    return torch.count_nonzero(red_clr_tensor, dim=1)         # [b]

# %% ../../src/dataset/balancing.ipynb #36032308-bd0e-4409-9db0-9d89fc258e5a
def add_balance_fn_quantile_qc_length(indices: Union[np.ndarray, torch.Tensor], 
                                      x: Union[np.ndarray, torch.Tensor], 
                                      y: Union[np.ndarray, torch.Tensor], 
                                      *z, 
                                      padding_token: int = 0,
                                      balance_quantile: float = 0.5,
                                      device: torch.device = torch.device("cpu"),
                                      quantile_length_weights: Optional[Callable[[torch.Tensor, torch.Tensor], torch.Tensor]] = None) -> torch.Tensor:
    """Balances according to gate length."""
    
    xb = x[indices].to(device)
    l  = get_tensor_gate_length(xb, padding_token=padding_token).to(device)
    
    l_uniques, l_uniques_cnt = torch.unique(l, dim=0, return_counts=True)

    #-----------------------------------
    # samples = torch.min(l_uniques_cnt)
    # samples = torch.median(l_uniques_cnt)
    samples = torch.quantile(l_uniques_cnt.float(), balance_quantile, interpolation='nearest', dim=0).to(l_uniques_cnt.dtype)
    samples = max(samples, 2)

    #-----------------------------------
    sub_ind = list()   
    for l_unique in l_uniques.to(device):      
        comp = (l==l_unique)
        ind  = comp.nonzero().squeeze().cpu()
        
        if ind.dim() > 0:
            if exists(quantile_length_weights):
                _samples = int(quantile_length_weights(l_unique, samples))
            else:
                _samples = samples
            
            ind = dahe.shuffle_tensor_dataset(ind) 
            ind = ind[:_samples]
        else:
            ind = ind[None]
        
        sub_ind.append(ind)

    sub_ind = torch.cat(sub_ind, dim=0)
      
    indices = indices[sub_ind]
    
    if indices.ndim < 1:  
        indices = indices[None]
    
    return indices
