"""Base class for managing, loading and saving."""

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

# %% auto #0
__all__ = ['ConfigDatasetConfig', 'ConfigDataset']

# %% ../../src/dataset/config_dataset.ipynb #06272f6f-b4e3-4504-a90a-feebbf6ad821
from ..imports import *
from ..utils.config_loader import *
from .dataset_helper import *

from huggingface_hub import snapshot_download

# %% ../../src/dataset/config_dataset.ipynb #a2688c32-b901-4ac1-9ed6-de0dbd12108b
@dataclass
class ConfigDatasetConfig:
    """Config `dataclass` used for storage."""
    store_dict: dict 
    dataset_to_gpu: bool

# %% ../../src/dataset/config_dataset.ipynb #36032308-bd0e-4409-9db0-9d89fc258e5a
class ConfigDataset():  
    """Base class for datasets, manages loading and saving."""
    
    req_params     = [f.name for f in dataclasses.fields(ConfigDatasetConfig)]
    comment        = ""
    add_balance_fn = None
    
    def __init__(self, device: torch.device=torch.device("cpu"), save_type=None, **parameters) -> None:
        self.save_type = default(save_type, "safetensors")
        
        req_params = self.req_params      
        for p in req_params:
            if p not in parameters: raise RuntimeError(f"Missing parameter `{p}` in argument `**parameters: dict`")           

        self.device = device      #parameters will overwrite passed device
        
        for k,v in parameters["store_dict"].items(): 
            if   v == "tensor"     : setattr(self, str(k), torch.tensor([0], device=self.device))
            elif v == "tensor_list": setattr(self, str(k), [torch.tensor([0], device=self.device)]) 
            elif v == "list"       : setattr(self, str(k), ["list str entry"]) 
            elif v == "numpy"      : setattr(self, str(k), np.array(["numpy str entry"]))                                             
            else                   : raise RuntimeError(f"Unknown type `{v}` in argument parameters[`store_dict`]")
                                             
        for k,v in parameters.items(): setattr(self, str(k), v)
           
    def to(self, device: torch.device, excepts=[], **kwargs):
        self.device = device
        
        for k,v in self.store_dict.items(): 
            if k in excepts: continue
            
            if v == "tensor": 
                x = getattr(self, str(k)).to(device, **kwargs)
                setattr(self, str(k), x)
                
            elif v == "tensor_list": 
                x = getattr(self, str(k))
                x = [ix.to(device, **kwargs) for ix in x]
                setattr(self, str(k), x)
  
        return self

    def memory_summary(self) -> None:
        print("##################### Dataset memory summary #####################")
        print("Name          || Type                       || Memory     || Device  || Shape")
        print("---------------------------------------------------------------")
        
        total_mem = 0.0
        byte_to_giga = 1 / (1024**3)
    
        for k,v in self.store_dict.items(): 
            mem = 0.0
            dev = "None"
            shape = "None"
            dtype = "None"
            
            x = getattr(self, str(k))
            
            if v == "tensor":
                mem += float(x.dtype.itemsize) * np.prod([s for s in x.shape], dtype=np.double) * byte_to_giga
                dev = x.device
                shape = x.shape
                dtype = x.dtype
            
            elif v == "tensor_list":             
                dev = []
                for x_i in x:
                    mem += float(x_i.dtype.itemsize) * np.prod([s for s in x_i.shape], dtype=np.double) * byte_to_giga
                    dev.append(x_i.device)
                shape = (len(x), x[0].shape)
                dtype = x[0].dtype
                    
            elif v == "list": 
                shape = (len(x))
                dtype = "python"
            
            elif v == "numpy": 
                shape = x.shape
                dtype = x.dtype

            
            print(f"  - [{str(k):>8}]  ({str(dtype):>15} {str(v):>8}):    {mem:3.4f} GB     ({str(dev):6})  | {shape}")
            total_mem += mem
            
        print("--------------------------------------")
        print(f"  Total memory used: {total_mem:3.4f} GB ")
        print("---------------------------------------------------------------")

    #----------------------------
    
    def x_y_preprocess(self, balance_max=None, shuffle=False, max_samples=None, make_unique=True):
        z_proc = []
        for k,v in self.store_dict.items(): 
            if k != "x" and k != "y":
                z_proc.append(getattr(self, k))
                              
        x_proc, y_proc = self.x, self.y
        
        #---------------------
        if shuffle:
            x_proc, y_proc, *z_proc = shuffle_tensor_dataset(x_proc, y_proc, *z_proc)
                
        if exists(max_samples):
            x_proc = x_proc[:max_samples]
            y_proc = y_proc[:max_samples]
            z_proc = (iz[:max_samples] for iz in z_proc)   
        
        #---------------------
        t = self.store_dict["y"]
        if exists(balance_max): 
            if t == "tensor" or t == "numpy": x_proc, y_proc, *z_proc = balance_tensor_dataset(x_proc, y_proc, *z_proc, make_unique=make_unique, shuffle_lables=shuffle, 
                                                                                               samples=balance_max, add_balance_fn=self.add_balance_fn, njobs=1) 
            else:                             print(f"[WARNING]: Unsupported y type: `{t}`. Not balancing dataset!")
        else: print(f"[INFO]: Not balancing dataset!  {balance_max=}")
          
        #---------------------
        if shuffle:
            x_proc, y_proc, *z_proc = shuffle_tensor_dataset(x_proc, y_proc, *z_proc)
            
        return x_proc, y_proc, *z_proc
    
    def valid_split(self, x, y, *z, p_valid=0.1, y_type=None, split_sequential=False):
        """
        split_sequential ... if true split data ordered (valid-train order), else split randomly (the same as shuffle and then seq. split)
        """
        
        if split_sequential: ind = torch.arange(x.shape[0])
        else:                ind = torch.randperm(x.shape[0]) 
       
        splits = max(int(x.shape[0] * p_valid), 1)  
        ind, ind_valid = ind[splits:], ind[:splits]

        #### Note: advanced indexing always creates copy not view. So we can skip the .clone()
        x, x_valid = x[ind], x[ind_valid]
        
        t = y_type if exists(y_type) else self.store_dict["y"]
        if   t == "tensor" : y, y_valid = y[ind], y[ind_valid]   
        elif t == "numpy":   y, y_valid = y[ind], y[ind_valid]
        
        z = list(z)
        z_valid = [None] * len(z)
        for i, iz in enumerate(z):
            # assert tensors for now
            z[i], z_valid[i] = iz[ind], iz[ind_valid]
            
        z, z_valid = tuple(z), tuple(z_valid)
               
        return x, x_valid, y, y_valid, (z, z_valid)
    
    def get_dataloaders(self, batch_size, p_valid=0.1, balance_max=None, max_samples=None, y_on_cpu=False, shuffle=True):
        #-------------------------
        # valid split and to device
        
        x_proc, y_proc, *z_proc              = self.x_y_preprocess(balance_max=balance_max, max_samples=max_samples, shuffle=shuffle)  
        x, x_valid, y, y_valid, (z, z_valid) = self.valid_split(x_proc, y_proc, *z_proc, p_valid=p_valid)


        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 self.params_config.dataset_to_gpu: 
            train_loader = DataLoader(dataset=ds      , batch_size=batch_size, shuffle=True)
            valid_loader = DataLoader(dataset=ds_valid, batch_size=batch_size, shuffle=True)

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

        self.dataloaders = DataLoaders(train_loader, valid_loader)        
        return self.dataloaders
    
    #----------------------------
    
    @property
    def params_config(self):
        params_config = {}              
        for p in self.req_params: params_config[p] = getattr(self, p)
            
        if type(self) == ConfigDataset:
            params_config = ConfigDatasetConfig(**params_config)
        return params_config   

    #----------------------------
    
    def get_config(self, save_path=None, without_metadata=False):
        if not without_metadata:       
            config = {}
            config["target"]         = class_to_str(type(self))
            config["device"]         = str(self.device)
            config["comment"]        = self.comment
            config["save_path"]      = self.save_path if hasattr(self, "save_path") and not exists(save_path) else save_path
            config["save_datetime"]  = datetime.now().strftime("%m/%d/%Y %H:%M:%S")
            config["save_type"]      = self.save_type
            config["params"]         = self.params_config  
        else:
            config = self.params_config  
        
        self.config = config        
        return config
            
    def save_dataset(self, config_path: str, save_path: str):
        if exists(config_path): os.makedirs(config_path[:config_path.rfind("/")] + "/", exist_ok=True)
        if exists(save_path):   os.makedirs(save_path[:save_path.rfind("/")] + "/", exist_ok=True)
        
        config = self.get_config(save_path, without_metadata=False)
        save_dict_yaml(config, config_path)                                   
        self.store_x_y(save_path)          
   
    #----------------------------

    def check_save_type(self, save_path):
        if exists(self.save_type) and exists(save_path):
            if not save_path.endswith(f".{self.save_type}"):
                save_path += f".{self.save_type}"
        return save_path
    
    def store_x_y(self, path_str):       
        for k,v in self.store_dict.items(): 
            x = getattr(self, str(k))

            # torch.save(x, path_str + f"_{k}.pt")
            store_tensor({"0": x}, self.check_save_type(path_str + f"_{k}"), type=v)
                           
    def load_x_y(self, path_str, device: Optional[torch.device] = None, make_contiguous: bool = True):
        self.save_path = path_str
        
        for k,v in self.store_dict.items():    
            # x = torch.load(path_str + f"_{k}.pt", map_location=device)
            x = load_tensor(self.check_save_type(path_str + f"_{k}"), device, type=v)

            if isinstance(x, dict):
                x = x["0"]

            if v == "tensor" and make_contiguous:
                x = x.contiguous() #load memmap into memory
            
            setattr(self, str(k), x)
        
    #----------------------------
    
    @staticmethod
    def from_config(config, device: torch.device, save_path: Optional[str] = None, make_contiguous: bool = True):
        """Use this if we have a loaded config."""
        
        config_dataset = instantiate_from_config(config)
        
        if "comment" in config: config_dataset.comment = config["comment"]
        
        #--------------------------------  

        config_dataset.save_type = config.pop("save_type", None)
        
        if not exists(save_path):            
            if "save_path" in config: save_path = config["save_path"]
            else:                     print("[INFO]: Found no key `save_path` path in config and no `save_path` arg provided.")
                                  
        if exists(save_path): config_dataset.load_x_y(save_path, device=device, make_contiguous=make_contiguous)
        else:                 print("[INFO]: No save_path` provided. Nothing loaded.")

        #--------------------------------
        
        config_dataset = config_dataset.to(device)
        print(f"[INFO]: Instantiated config_dataset from given config on {device}.")
        
        return config_dataset
    
    @classmethod
    def from_config_file(cls, config_path, device: torch.device, save_path:  Optional[str] = None, make_contiguous: bool = True):
        """
        Load a dataset from a config file. 
        If this method is called with `ConfigDataset.from_config_file` we use the given `target`, else use the caller class.
        """
        config = load_config(config_path)
        if cls is not ConfigDataset:
            config["target"] = class_to_str(cls)      
        return cls.from_config(config, device, save_path, make_contiguous)

    @classmethod
    def from_huggingface(cls, repo_id: str, device: torch.device, **kwargs):  
        """Load a dataset directly from Huggingface."""
        dataset_path = snapshot_download(repo_id=repo_id, repo_type="dataset", allow_patterns=["*.pt", "*.yaml", "*.safetensors"], **kwargs) 

        try:
            name = repo_id.split("/")[-1]
            dataset = cls.from_config_file(config_path=dataset_path+f"/{name}.yaml", device=device, save_path=dataset_path+f"/{name}")  
        except Exception as e:
            dataset = cls.from_config_file(config_path=dataset_path+"/config.yaml", device=device, save_path=dataset_path+"/dataset")  
        
        return dataset 
