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

# %% auto 0
__all__ = ['Cached_OpenClip_Dataset']

# %% ../../src/dataset/cached_qc_dataset.ipynb 3
from ..imports import *
from .qc_dataset import Qc_Config_Dataset
from .config_dataset import Config_Dataset
from ..config_loader import *

# %% ../../src/dataset/cached_qc_dataset.ipynb 4
class Cached_OpenClip_Dataset(Qc_Config_Dataset):
    """Adds `.caching` to the `Quantum circuit dataset` class."""
    
    def x_y_preprocess(self, balance_max, max_samples=None):
        x_proc, y_proc, *z = super().x_y_preprocess(balance_max=balance_max, max_samples=max_samples)        
        y_proc = self.caching(y_proc)
        return x_proc, y_proc, *z
    
    def caching(self, y_proc, y_on_cpu=False):
        print("[INFO]: Generate cache: converting tensors to str and tokenize")   
        
        print(" - to str list")  
        if isinstance(y_proc, (torch.Tensor, torch.IntTensor, torch.FloatTensor, torch.LongTensor)):         
            y_str = [str(i) for i in y_proc.cpu().tolist()]
        elif isinstance(y_proc, list): 
            y_str = []
            for iy in y_proc:                
                if isinstance(iy, np.ndarray): y_str += [str(i) for i in iy.tolist()]        # list of numpy arrays
                else:                          y_str += [str(i) for i in iy.cpu().tolist()]  # list of tensors
        elif isinstance(y_proc, np.ndarray):
            y_str = [str(i) for i in y_proc.tolist()]
            
        else: raise NotImplementedError()
                            
        print(" - tokenize_and_push_to_device")  
        y_tok = self.text_encoder.tokenize_and_push_to_device(y_str, to_device= not y_on_cpu)
        if y_on_cpu: y_tok = y_tok.cpu()
        
        
        #now for using cache we need the uniques and the corresponding indices of the uniques
        y_uniques, y_ptrs  = torch.unique(torch.cat([self.text_encoder.empty_token.to(y_tok.device), y_tok]), dim=0, return_inverse=True)
    
        cached_empty_token_index = y_ptrs[0]  #store what index the empty token has   
        y_ptrs                   = y_ptrs[1:] #remove the cat empty token
      
        #use cache
        print(" - generate_cache")  
        self.text_encoder.generate_cache(tokens=y_uniques, cached_empty_token_index=cached_empty_token_index, y_on_cpu=y_on_cpu)
      
        print("[INFO]: Generated cache")  
        return y_ptrs
    
    #-------------------------------------------
    
    def get_dataloaders(self, batch_size, text_encoder, p_valid=0.1, balance_max=None, max_samples=None):
        self.text_encoder = text_encoder    
        return super().get_dataloaders(batch_size, p_valid, balance_max, max_samples)
       
    #-------------------------------------------
        
    @staticmethod
    def from_config_file(config_path, device: torch.device, save_path: str=None):
        config = load_config(config_path)
        config["target"] = class_to_str(Cached_OpenClip_Dataset)               
        return Config_Dataset.from_config(config, device, save_path)
