"""Dataset that combines and handles multiple cached datasets."""

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

# %% auto #0
__all__ = ['MixedCachedOpenCLIPDatasetConfig', 'MixedCachedOpenCLIPDataset']

# %% ../../src/dataset/mixed_cached_dataset.ipynb #06272f6f-b4e3-4504-a90a-feebbf6ad821
from ..imports import *
from .cached_dataset import CachedOpenCLIPDataset, CachedOpenCLIPDatasetConfig, ConfigDataset
from .dataset_helper import *
from ..utils.misc_utils import DataLoaders, MemoryCleaner
from tensordict import TensorDict

# %% ../../src/dataset/mixed_cached_dataset.ipynb #119077c9-999b-44f7-8099-79037503d7e7
@dataclass
class MixedCachedOpenCLIPDatasetConfig(CachedOpenCLIPDatasetConfig):
    pad_constant: int
    collate_fn: str
    bucket_batch_size: int
    model_scale_factor: int

# %% ../../src/dataset/mixed_cached_dataset.ipynb #0037efb5-d3a9-46e4-94d1-3dd80297e934
class MixedCachedOpenCLIPDataset(CachedOpenCLIPDataset):  
    """Dataset that uses multiple cached dataset and combines them with padding, either i) Bucket or ii) Max."""
    
    req_params = [f.name for f in dataclasses.fields(MixedCachedOpenCLIPDatasetConfig)]

    #-----------------------------------
    @property
    def params_config(self):
        params_config = super().params_config  
        if type(self) == MixedCachedOpenCLIPDataset:
            params_config = MixedCachedOpenCLIPDatasetConfig(**params_config)
        return params_config  
        
    #-----------------------------------
    # functions to combine multiple datasets together

    @classmethod
    def _preprocess_datasets(dataset_cls, datasets, device, balance_maxes, max_samples, shuffle, 
                             make_unique, pad_constant, model_scale_factor, parameters, **kwargs):
        xs = []
        ys = []
        zs = []
        cs = []
        
        if isinstance(max_samples, int):
            max_samples = [max_samples] * len(datasets)
        else:
            assert isinstance(max_samples, (list, np.ndarray))

        if isinstance(balance_maxes, int):
            balance_maxes = [balance_maxes] * len(datasets)
        else:
            assert isinstance(balance_maxes, (list, np.ndarray))
        
        for i, (dataset, balance_max) in tqdm(enumerate(zip(datasets, balance_maxes)), total=len(datasets)):

            x, y, z, *c = dataset_cls._preprocess_dataset(dataset, device, balance_max, max_samples, i, shuffle, make_unique, pad_constant, model_scale_factor, parameters, **kwargs)
            MemoryCleaner.purge_mem()
            
            #combine datasets
            xs.append(x.cpu())  
            ys.append(y)
            zs.append(z.cpu()) 
            cs.append([ic.cpu() for ic in c])

            del x
            del y
            del z
            del c
            
            for k in datasets[i].store_dict.keys(): 
                setattr(datasets[i], str(k), None)
            del dataset
            
            MemoryCleaner.purge_mem()

        return xs, ys, zs, cs
      
    @staticmethod
    def _add_missing_conditions(parameters, dataset, c, batch_size, device):
        # if c is missing something of the union we set it to a zero tensor, e.g. used for combining SRV with compilation
        c_temp = []
        c_temp_index = 0
        
        for k,v in parameters["store_dict"].items(): 
            if k != "x" and k != "y" and k != "z":   
                if k not in dataset.params_config.store_dict:
                    empty_tensor = torch.zeros((1,), device=device)
                    
                    if k == "U": #scetchy hardcoded for compilation
                        empty_tensor = torch.zeros((batch_size, 2, 1, 1), device=device) # unitary is [b, Re/Im, 2^n, 2^n]
                    
                    c_temp.append(empty_tensor) 
                    
                else: # done to conserve the ordering of c args!!!
                    c_temp.append(c[c_temp_index])
                    c_temp_index += 1

        return c_temp

    @staticmethod
    def _reorder_to_buckets(parameters, bucket_batch_size, xs, ys, zs, cs):
        for i, (xi,yi,zi, ci) in enumerate(zip(xs, ys, zs, cs)):  #cut rest of batch        
            b_mult = int(np.floor(xi.shape[0] / bucket_batch_size) * bucket_batch_size)  
            
            xs[i] = xi[None, :b_mult].reshape((b_mult//bucket_batch_size, bucket_batch_size, *xi.shape[1:])) 
            zs[i] = zi[None, :b_mult].reshape((b_mult//bucket_batch_size, bucket_batch_size, *zi.shape[1:]))
            
            v = parameters["store_dict"]["y"]
            if v == "tensor" or v == "numpy": 
                ys[i] = yi[None, :b_mult].reshape((b_mult//bucket_batch_size, bucket_batch_size, *yi.shape[1:]))    
            else: raise NotImplementedError("")
        
            #----
            #For U, etc
            add_ind = 0
            for k,v in parameters["store_dict"].items(): 
                if k != "x" and k != "y" and k != "z":                             
                    if v == "tensor" or v == "numpy": 
                        cs[i][add_ind] = ci[add_ind][None, :b_mult].reshape((b_mult//bucket_batch_size, bucket_batch_size, *ci[add_ind].shape[1:]))   
                    else: raise NotImplementedError("")                      
                    add_ind += 1  

        return xs, ys, zs, cs

    @staticmethod
    def _pad_conditions(parameters, bucket_batch_size, c, unitary_pad=None, params_pad=None, pad_with_memmap=False):
        ci_list   = []
        ci_k_list = []

        memmap_cleans = [] #TensorDicts and paths we need to delete later
        
        def _alloc_mem(shape, k, c0_add_ind):
            # allocating zeros is better memory wise than torch.cat(ci_s) and F.pad(ci, pad, "constant", 0)
            mem = np.prod(shape) * c0_add_ind.element_size() / (1024*1024*1024)
            print(f"[INFO]: allocate memory for {k} {shape} on {c0_add_ind.device} approx. {mem:.3f} GB")

            if pad_with_memmap:
                prefix_path = f"tmp_DELETE_pad_conditions_MixedCachedOpenCLIPDataset_{k}"
                print(f"[INFO]: (MixedCachedOpenCLIPDataset._pad_conditions): {pad_with_memmap=} allocating TensorDict using memmap_like at {prefix_path}")
                
                b, *_ = shape
                tensor_dict = TensorDict({"ci_s": torch.empty(shape, dtype=c0_add_ind.dtype),
                                         }, batch_size=[b])
                tensor_dict = tensor_dict.memmap_like(prefix=prefix_path)
                        
                ci_s = tensor_dict["ci_s"]
                memmap_cleans.append((tensor_dict, prefix_path))
            else:
                ci_s = torch.zeros(shape, device=c0_add_ind.device, dtype=c0_add_ind.dtype)    
                
            return ci_s

        add_ind = 0
        for k,v in parameters["store_dict"].items(): 
            if k != "x" and k != "y" and k != "z":   
                                     
                if v == "tensor" and k == "U":    # hardcoded U padding !!
                    assert exists(unitary_pad) and isinstance(unitary_pad, int)
        
                    n = sum([ci[add_ind].shape[0] for ci in c])
                    if bucket_batch_size > 0: shape = (n, bucket_batch_size, 2, unitary_pad, unitary_pad)
                    else:                     shape = (n,                    2, unitary_pad, unitary_pad)
                            
                    ci_s = _alloc_mem(shape, k, c[0][add_ind])            

                    #tensor product pad, else was zero pad
                    if 1:
                        run_i = 0
                        for i,ci in enumerate(c):
                            ci = ci[add_ind]     

                            assert ci.shape[-2]==ci.shape[-1]
                            U_side = ci.shape[-2]
                            for jj in range(unitary_pad//U_side):                             
                                ci_s[run_i:run_i+ci.shape[0], ..., U_side*jj:U_side*(jj+1), U_side*jj:U_side*(jj+1)] = ci.to(ci_s.device)                          
                            
                            run_i += ci.shape[0]
            
                        ci_list.append(ci_s)
                        ci_k_list.append(k)
            
                        add_ind += 1
                        continue
                
                elif v == "tensor" and k == "params": # hardcoded paramter padding !!
                    assert exists(params_pad) #and len(list(params_pad))==2
                                       
                    n = sum(ci[add_ind].shape[0] for ci in c)
                    if bucket_batch_size > 0: shape = (n, bucket_batch_size, *params_pad)
                    else:                     shape = (n,                    *params_pad)
    
                    ci_s = _alloc_mem(shape, k, c[0][add_ind])    
                                                     
                elif v == "numpy": raise NotImplementedError("")   
                else:              raise NotImplementedError("")           
    
                
                run_i = 0
                for i,ci in enumerate(c):
                    ci = ci[add_ind]     
                    ci_s[run_i:run_i+ci.shape[0], ..., :ci.shape[-2], :ci.shape[-1]] = ci                          
                    run_i += ci.shape[0]

                ci_list.append(ci_s)
                ci_k_list.append(k)
    
                add_ind += 1

        return ci_list, ci_k_list, memmap_cleans
    
    @classmethod
    def _create_train_valid_datasets(dataset_cls, device, parameters, test_split, x, y, z, ci_list, ci_k_list, shuffle: bool = True):
        splits = max(int(x.shape[0] * test_split), 1)

        if shuffle:
            x, y, z, *ci_list = shuffle_tensor_dataset(x, y, z, *ci_list)

        x, x_test =  x[splits:], x[:splits]
        y, y_test =  y[splits:], y[:splits]
        z, z_test =  z[splits:], z[:splits]

        print(f"Split: Train {x.shape[0]} - Test {x_test.shape[0]} \n")
        
        dataset = dataset_cls(device, **parameters)        
        dataset.x = x
        dataset.y = y
        dataset.z = z
      
        dataset_test = dataset_cls(device, **parameters)        
        dataset_test.x = x_test
        dataset_test.y = y_test
        dataset_test.z = z_test
        
        for ci, k in zip(ci_list, ci_k_list):       
            ci, ci_test =  ci[splits:], ci[:splits]
            
            setattr(dataset     , str(k), ci)
            setattr(dataset_test, str(k), ci_test)
        
        return dataset, dataset_test
    
    #-----------------------------------
    
    def get_dataloaders(self, batch_size, text_encoder, p_valid=0.1, y_on_cpu=False, return_tensor_datasets=False, shuffle=True, shuffle_cpu_copy=True, caching=True):
        #-------------------------
        # caching
        
        self.text_encoder = text_encoder

        print("[DEBUG]: run get_dataloaders.x_y_preprocess", flush=True)
        x_proc, y_proc, *z_proc = ConfigDataset.x_y_preprocess(self, 
                                                               balance_max=None, 
                                                               shuffle=False, 
                                                               max_samples=None, 
                                                               make_unique=False)    # ... z_proc is `'z' and all other 'c'
        if caching:
            if self.bucket_batch_size <= 0:        
                y_proc = self.caching(y_proc, y_on_cpu=y_on_cpu)
                          
            else:                    
                y_proc = self.caching([yi.reshape((-1)) for yi in y_proc], y_on_cpu=y_on_cpu)
                y_proc = y_proc.reshape((-1, self.bucket_batch_size))
        
        #-------------------------
        # valid split and to device

        print("[DEBUG]: run get_dataloaders.valid_split", flush=True)
        x, x_valid, y, y_valid, (z, z_valid) = self.valid_split(x_proc, y_proc, *z_proc, p_valid=p_valid, y_type="tensor", split_sequential=False)

        if self.params_config.dataset_to_gpu:
            x, x_valid = x.to("cuda"), x_valid.to("cuda")
            z, z_valid = list(iz.to("cuda") for iz in z), list(iz_valid.to("cuda") for iz_valid in z_valid)

            if not y_on_cpu:
                y, y_valid = y.to("cuda"), y_valid.to("cuda")

        #-------------------------
        # create dataloaders

        ds       = TensorDataset(x, y, *z)
        ds_valid = TensorDataset(x_valid, y_valid, *z_valid)

        if return_tensor_datasets:
            return ds, ds_valid

        if isinstance(self.collate_fn, str):
            collate_fn = getattr(self, self.collate_fn, None)
        else:
            collate_fn = self.collate_fn
        
        if not exists(collate_fn):
            print("[WARNING]: self.collate_fn does not exist, using torch.utils.data.default_collate.")
            collate_fn = torch.utils.data.default_collate

        if self.params_config.dataset_to_gpu: 
            train_loader = DataLoader(dataset=ds      , batch_size=batch_size, shuffle=True, collate_fn=collate_fn)
            valid_loader = DataLoader(dataset=ds_valid, batch_size=batch_size, shuffle=True, collate_fn=collate_fn)

        else:              
            train_loader = DataLoader(dataset=ds      , batch_size=batch_size, shuffle=True, pin_memory=True, num_workers=4, collate_fn=collate_fn)
            valid_loader = DataLoader(dataset=ds_valid, batch_size=batch_size, shuffle=True, pin_memory=True, num_workers=4, collate_fn=collate_fn)

        self.dataloaders = DataLoaders(train_loader, valid_loader)        
        return self.dataloaders
    
    #-----------------------------------
    
    @staticmethod
    def from_datasets(*args, **kwargs):
        raise NotImplementedError()
