"""Model base class that handles loading and storing from/to config files."""

# AUTOGENERATED! DO NOT EDIT! File to edit: ../../src/models/config_model.ipynb.

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

# %% ../../src/models/config_model.ipynb #7db9efd6-abbc-4256-b0d9-7f7da64d81b4
from ..imports import *
from ..utils.config_loader import *
from ..utils.misc_utils import *
from datetime import datetime

# %% ../../src/models/config_model.ipynb #90e32385-485c-4bb0-863a-9b6390e8b9df
class ConfigModel(nn.Module):
    """A basic `nn.Module` with IO functionality."""
    
    def __init__(self, save_type=None): 
        super().__init__()
        self.save_type = default(save_type, "safetensors")
    
    def freeze(self, freeze: bool = True):
        if freeze: self.eval()
        else:      self.train()

        for param in self.parameters(): 
            param.requires_grad = not freeze 

        #Todo: add a debose/debug log here
    
    def unfreeze(self):
        self.freeze(False)
    
    #---------------------

    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 get_config(self, save_path=None, without_metadata=False):
        if not without_metadata:       
            config = {}
            config["target"]         = class_to_str(type(self)) 
            config["save_path"]      = self.check_save_type(self.save_path) if hasattr(self, "save_path") and not exists(save_path) else self.check_save_type(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 store_model(self, config_path: str=None, save_path: str=None, without_metadata=False):        
        
        config = self.get_config(save_path, without_metadata)
    
        if exists(config_path):
            if without_metadata: save_dataclass_yaml(config, config_path)
            else               : save_dict_yaml(config, config_path)            
        
        if exists(save_path):
            store_model_state_dict(self.state_dict(), self.check_save_type(save_path)) 
    
    #---------------------
    
    @staticmethod
    def from_config(config, device: torch.device, save_path: str=None, verbose=True, silent=False, freeze: Optional[bool] = None):  
        """Use this if we have a loaded config. Maybe within other classes (e.g. pipeline and nested models)"""

        _config = copy.deepcopy(config)
        
        if exists(device): _config["device"] = device   # for loading sub-models
        else:              device = _config.pop("device", "cpu")

        if exists(freeze):
            _freeze = freeze
      
        else:
            if "is_frozen" in _config: 
                _freeze = _config.pop("is_frozen", None)
                if not_exists(_freeze):
                    raise RuntimeError(f"The `is_frozen` flag in `config` is invalid. Please provide a boolean. `is_frozen` is: {freeze}")
            else:
                _freeze = True
                #print(f"[INFO]: `{class_to_str(type(model))}`. No valid `is_frozen` flag in `config`. Model is frozen by default.")
        
        #--------------------------------
        # instantiate model
        model = instantiate_from_config(_config)
        model = model.to(device) 
        if not silent: print(f"[INFO]: `{class_to_str(type(model))}` instantiated from given `config` on {device}.")

        #--------------------------------        
        # load pretrained weights

        model.save_type = _config.pop("save_type", None)

        if exists(model.save_type):
            if not exists(save_path):            
                if "save_path" in _config: 
                    save_path = model.check_save_type(_config["save_path"])
    
                                      
            if exists(save_path):        
                state_dict = load_model_state_dict(model.check_save_type(save_path), device)
                
                m, u = model.load_state_dict(state_dict, strict=False)
         
                if len(m) + len(u) > 0 and verbose:
                    print(f"[WARNING]: missing keys: {m}")
                    print(f"[WARNING]: unexpected keys: {u}")
    
            else:
                if not silent: print(f"[INFO]: `{class_to_str(type(model))}`. No `save_path` provided. Found no key `save_path` in `config`. No state dict loaded.")
        else:
            if not silent: print(f"[INFO]: `{class_to_str(type(model))}`. Found no key `save_type` in `config`. No state dict loaded.")
        
        #--------------------------------
        # freeze    

        if exists(_freeze):
            model.freeze(_freeze)
            if not silent: print(f"[INFO]: `{class_to_str(type(model))}`. Freeze model: {_freeze}")
        else:
            if not silent: print(f"[INFO]: `{class_to_str(type(model))}`. No valid `is_frozen` flag in `config`. Model is frozen by default.")

        return model
    
    @staticmethod
    def from_config_file(config_path, device: torch.device, save_path: str=None):    
        config = load_config(config_path)
        return ConfigModel.from_config(config, device, save_path)       
