"""Dataset for quantum circuits."""

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

# %% auto #0
__all__ = ['CircuitsConfigDatasetConfig', 'CircuitsConfigDataset', 'MixedCircuitsConfigDatasetConfig',
           'MixedCircuitsConfigDataset']

# %% ../../src/dataset/circuits_dataset.ipynb #06272f6f-b4e3-4504-a90a-feebbf6ad821
from ..imports import *
from .cached_dataset import CachedOpenCLIPDataset, CachedOpenCLIPDatasetConfig
from .mixed_cached_dataset import MixedCachedOpenCLIPDataset, MixedCachedOpenCLIPDatasetConfig
from ..utils.config_loader import *
from .config_dataset import ConfigDataset
from .dataset_helper import shuffle_tensor_dataset
from ..utils.misc_utils import MemoryCleaner

# %% ../../src/dataset/circuits_dataset.ipynb #016fc327-f986-4d69-b5f0-1b39466fb528
@dataclass
class CircuitsConfigDatasetConfig(CachedOpenCLIPDatasetConfig):
    optimized: bool
    random_samples: int  
    num_of_qubits: int  
    min_gates: int 
    max_gates: int 
    max_params: int
    gate_pool: list[str]

# %% ../../src/dataset/circuits_dataset.ipynb #36032308-bd0e-4409-9db0-9d89fc258e5a
class CircuitsConfigDataset(CachedOpenCLIPDataset):
    """Dataset for quantum circuits, access `gate_pool` directly and all other paras with `.params_config`"""
    
    req_params = [f.name for f in dataclasses.fields(CircuitsConfigDatasetConfig)]

    #-----------------------------------
    def __init__(self, device: torch.device=torch.device("cpu"), **parameters) -> None:
        super().__init__(device, **parameters)        

        
        if isinstance(list(parameters["gate_pool"])[0], str):
            self.gate_pool = list(parameters["gate_pool"])
            
        else:
            try:
                self.gate_pool = [get_obj_from_str(node) for node in parameters["gate_pool"]]   
            except Exception as er:
                print(f"[WARNING]: error => {er}")
                print(f"[WARNING]: gate_pool is passed as str")
                self.gate_pool = [str(node) for node in parameters["gate_pool"]]   
    
    @property
    def params_config(self):
        params_config = super().params_config         
        
        if type(self) == CircuitsConfigDataset:
            params_config = CircuitsConfigDatasetConfig(**params_config)
        return params_config   

# %% ../../src/dataset/circuits_dataset.ipynb #75c9a200-f9eb-42f9-b3c3-e074e377737a
@dataclass
class MixedCircuitsConfigDatasetConfig(CircuitsConfigDatasetConfig, MixedCachedOpenCLIPDatasetConfig):
    pass

# %% ../../src/dataset/circuits_dataset.ipynb #e804a8d7-dcf3-40e4-83a5-cd98207c8dea
class MixedCircuitsConfigDataset(CircuitsConfigDataset, MixedCachedOpenCLIPDataset):
    """
    Dataset that uses multiple cached dataset and combines them with padding, either i) Bucket or ii) Max.
    Also provides a corresponding `collate_fn` for training.
    """

    req_params = [f.name for f in dataclasses.fields(MixedCircuitsConfigDatasetConfig)]

    #-----------------------------------
    
    @property
    def params_config(self):
        params_config = super().params_config            
        if type(self) == MixedCircuitsConfigDataset:
            params_config = MixedCircuitsConfigDatasetConfig(**params_config)
        return params_config   

    #-----------------------------------

    def _get_cut_sizes(self, z):
        z_0 = torch.max(z[:, 0]) # space
        z_1 = torch.max(z[:, 1]) # time
        z_1 = (torch.ceil(z_1 / self.model_scale_factor) * self.model_scale_factor).to(torch.int32)
        return z_0, z_1
  
    def _cut(self, x, y, z):     
        z_0, z_1 = self._get_cut_sizes(z)
        
        x = x[:, :z_0, :z_1] # cut down to max [b, bits, time] of batch
        return x, y
    
    def _cut_compilation_params(self, x, y, p, U, z):  
        z_0, z_1 = self._get_cut_sizes(z)
        bit_exp  = 2**z_0
        
        x = x[:, :z_0, :z_1]              # cut down to max [b, bits, time] of batch
        p = p[:,    :, :z_1]              # cut down to max [b, nP  , time] of batch       
        U = U[:, :, :bit_exp, :bit_exp]   # [b, Re/Im, 2^n, 2^n]
        return x, y, p, U
    
    #-----------------------------------
    # BUCKET PADDING, all x,y are already passed as batch
        
    def cut_padding_Bucket_collate_fn(self, b):     
        """this function is called for training for every batch, order in b is store dict"""       

        x, y, z = b[0]
        x, y = self._cut(x, y, z)
        return x, y 

        
    def cut_padding_Bucket_collate_fn_compilation(self, b):     
        """this function is called for training for every batch"""    
        raise NotImplementedError()


    def cut_padding_Bucket_collate_fn_compilation_params(self, b):     
        """this function is called for training for every batch, order in b is store dict"""    
        
        b = b[0] # {'x': 'tensor', 'y': 'numpy', 'params': 'tensor', 'U': 'tensor', 'z': 'tensor'}
        
        x = b[0]
        y = b[1]  
        p = b[2]
        U = b[3]
        z = b[4]
        
        #---------------
        
        x, y, p, U = self._cut_compilation_params(x, y, p, U, z)
               
        return x, y, p, U
    
    #-----------------------------------
    # MAX PADDING, x are passes as sampled list (batch), std collate them
    
    def cut_padding_collate_fn(self, b):     
        """this function is called for training for every batch"""    
        x, y, z = torch.utils.data.default_collate(b)
        x, y    = self._cut(x, y, z)
        return x, y 

    def cut_padding_collate_fn_compilation(self, b):
        """this function is called for training for every batch"""    
        raise NotImplementedError()
        
    def cut_padding_collate_fn_compilation_params(self, b):
        """this function is called for training for every batch, order in b is store dict"""    
        # {'x': 'tensor', 'y': 'numpy', 'params': 'tensor', 'U': 'tensor', 'z': 'tensor'}
        x, y, p, U, z = torch.utils.data.default_collate(b)
        x, y, p, U    = self._cut_compilation_params(x, y, p, U, z)        
        return x, y, p, U   
        
    #-----------------------------------
    
    @staticmethod
    def _preprocess_dataset(dataset, device, balance_max, max_samples, i, shuffle, make_unique, pad_constant, 
                            model_scale_factor, parameters, max_gates, max_qubits):

        dataset = dataset.to(device)

        existing_z_type = dataset.store_dict.pop("z", None)  # remove z, as it would mess up `ConfigDataset.x_y_preprocess`, it would be put in `*c`.
        if exists(existing_z_type):
            assert existing_z_type == "tensor"
            z = dataset.z
        else:
            z = None
        
        x, y, *c = ConfigDataset.x_y_preprocess(dataset, balance_max=balance_max, max_samples=max_samples[i], shuffle=shuffle, make_unique=make_unique)       
        x = x.to(device)    # [b, s, t]   
        
        print(f" - dataset size after balancing {x.shape[0]}")

        #-------
        # store original size
        if not_exists(z):
            z = torch.zeros((x.shape[0], 2), device=device, dtype=torch.int32)
            z[:, 0] = max(dataset.params_config.num_of_qubits, 1)
            
            red_x   = torch.sum(x.abs(), dim=1)          # [b, t]   .. collaps the zeros to get circuit length
            z[:, 1] = torch.count_nonzero(red_x, dim=1)  # [b]         
            z[z[:, 1]==0, 1] = 1            

        # Create masks for space and time padding
        space_mask = torch.arange(x.shape[1], device=x.device).unsqueeze(0) >= z[:, 0].unsqueeze(1)
        time_mask  = torch.arange(x.shape[2], device=x.device).unsqueeze(0) >= z[:, 1].unsqueeze(1)

        # Apply masks to pad_constant to handle both dimensions
        x = torch.where(space_mask.unsqueeze(2), pad_constant, x)
        x = torch.where( time_mask.unsqueeze(1), pad_constant, x)
        
        z[:, 1] = (torch.ceil(z[:, 1] / model_scale_factor) * model_scale_factor).to(torch.int32) #for cut needs multiple

        #-------
        
        # now pad x, padding is defined from last dim forward!        
        pad = (0, max_gates-dataset.params_config.max_gates, 0, max_qubits-dataset.params_config.num_of_qubits) 
        x   = F.pad(x, pad, "constant", pad_constant)
        
        #-------

        c = MixedCachedOpenCLIPDataset._add_missing_conditions(parameters, dataset, c, x.shape[0], "cpu")

        dataset = dataset.to("cpu") #helps with gpu mem overflowing
        del dataset
            
        return x.cpu(), y, z.cpu(), *[ic.cpu() for ic in c]
    
    @staticmethod
    def from_datasets(datasets: list[CircuitsConfigDataset], balance_maxes: list, pad_constant, device: torch.device=torch.device("cpu"), bucket_batch_size=None, 
                      max_samples=None, shuffle=True, make_unique=True, test_split=0.05, pad_with_memmap=False, **parameters):
        if pad_constant == 0:
            print("[WARNING]: >pad_constant == 0<; This could be an error!")
        
        model_scale_factor = parameters["model_scale_factor"]
        
        max_qubits  = max(dataset.params_config.num_of_qubits for dataset in datasets)
        max_gates   = max(dataset.params_config.max_gates     for dataset in datasets)
        max_gates   = int(np.ceil(max_gates /model_scale_factor) * model_scale_factor)
        max_params  = max(dataset.params_config.max_params for dataset in datasets)
        
        parameters["num_of_qubits"]     = max_qubits
        parameters["max_gates"]         = max_gates
        parameters["max_params"]        = max_params
        parameters["random_samples"]    = sum([dataset.params_config.random_samples for dataset in datasets])
        parameters["min_gates"]         = min([dataset.params_config.min_gates      for dataset in datasets])
        parameters["comment"]           = f"Generated with 'from_datasets' with {len(datasets)} datasets. Qubits: {[dataset.params_config.num_of_qubits for dataset in datasets]}."
        parameters["pad_constant"]      = pad_constant
        parameters["bucket_batch_size"] = bucket_batch_size
         
        parameters["store_dict"] = {}
        for dataset in datasets:
            parameters["store_dict"] |= dataset.params_config.store_dict   #needs python 3.9 for union of dict  
        parameters["store_dict"]["z"]   = "tensor" #add special item

        #-----------------
        
        xs, ys, zs, cs = MixedCircuitsConfigDataset._preprocess_datasets(datasets, device, balance_maxes, max_samples, shuffle, make_unique, pad_constant, 
                                                                         model_scale_factor, parameters, max_gates=max_gates, max_qubits=max_qubits)            
        #-----------------

        has_U = "U" in parameters["store_dict"]
        has_p = "params" in parameters["store_dict"]
        
        if bucket_batch_size > 0:
            collate_fn_name = MixedCircuitsConfigDataset.cut_padding_Bucket_collate_fn.__name__
            if has_U: 
                collate_fn_name = MixedCircuitsConfigDataset.cut_padding_Bucket_collate_fn_compilation.__name__
                if has_p: 
                    collate_fn_name = MixedCircuitsConfigDataset.cut_padding_Bucket_collate_fn_compilation_params.__name__
        
        else:
            collate_fn_name = MixedCircuitsConfigDataset.cut_padding_collate_fn.__name__   
            if has_U: 
                collate_fn_name = MixedCircuitsConfigDataset.cut_padding_collate_fn_compilation.__name__
                if has_p: 
                    collate_fn_name = MixedCircuitsConfigDataset.cut_padding_collate_fn_compilation_params.__name__

        parameters["collate_fn"] = collate_fn_name
        
        #-----------------
        if bucket_batch_size > 0:
            xs, ys, zs, cs = MixedCachedOpenCLIPDataset._reorder_to_buckets(parameters, bucket_batch_size, xs, ys, zs, cs)
                                     
        x = torch.cat(xs)
        y = ys                 # torch.cat(ys) is wrong,  y is list of numpy or str!! not a tensor
 
        if isinstance(y, list): 
            match parameters["store_dict"]["y"]:
                case "numpy":  y = np.concatenate(y, axis=0)
                case "tensor": y = torch.cat(y, dim=0)
                case _: raise NotImplementedError()
            
        z = torch.cat(zs)
        c = cs
        
        #-----------------

        params_pad  = (max_params, max_gates)
        unitary_pad = 2**max_qubits
        
        ci_list, ci_k_list, memmap_cleans = MixedCachedOpenCLIPDataset._pad_conditions(parameters, bucket_batch_size, c, unitary_pad=unitary_pad, params_pad=params_pad, pad_with_memmap=pad_with_memmap)
        
        #-----------------    

        mixed_CircuitsConfigDataset, mixed_CircuitsConfigDataset_test = \
                MixedCircuitsConfigDataset._create_train_valid_datasets(device, parameters, test_split, x, y, z, ci_list, ci_k_list, shuffle=shuffle)

        if pad_with_memmap:
            mixed_CircuitsConfigDataset.memmap_cleans      = memmap_cleans
            mixed_CircuitsConfigDataset_test.memmap_cleans = memmap_cleans
        
        return mixed_CircuitsConfigDataset, mixed_CircuitsConfigDataset_test
