"""Functions to load and store models and datasets."""

# AUTOGENERATED! DO NOT EDIT! File to edit: ../../src/utils/config_loader.ipynb.

# %% auto #0
__all__ = ['class_to_str', 'load_config', 'config_to_dict', 'save_dataclass_yaml', 'save_dict_yaml', 'get_obj_from_str',
           'instantiate_from_config', 'store_model_state_dict', 'load_model_state_dict', 'store_tensor', 'load_tensor']

# %% ../../src/utils/config_loader.ipynb #06272f6f-b4e3-4504-a90a-feebbf6ad821
from ..imports import *
from omegaconf import OmegaConf

from safetensors.torch import load_file as load_safetensors
from safetensors.torch import save_file as save_safetensors
from safetensors.numpy import load_file as load_safetensors_np
from safetensors.numpy import save_file as save_safetensors_np
from safetensors import safe_open

# %% ../../src/utils/config_loader.ipynb #cb173637-3d18-4f94-8b95-76cda4117b1e
def class_to_str(cls):
    return str(cls)[8:-2]

# %% ../../src/utils/config_loader.ipynb #e0d5bc35-cc53-42fb-8fcd-8f2bc66c7c9b
def load_config(file_path):
    return OmegaConf.load(f"{file_path}")

# %% ../../src/utils/config_loader.ipynb #b105c04a-66d1-4450-8ee0-87aae618e60a
def config_to_dict(config):
    return OmegaConf.to_container(config)

# %% ../../src/utils/config_loader.ipynb #e3d81e5c-cf3d-4152-ab66-acd6e42ec3c9
def save_dataclass_yaml(data_obj, file_path):
    conf = OmegaConf.structured(data_obj)
    with open(file_path, 'w') as f:
        OmegaConf.save(config=conf, f=f)

# %% ../../src/utils/config_loader.ipynb #bf89bbfc-9d4e-442f-96ea-db1ab99505e9
def save_dict_yaml(dict_obj, file_path):
    conf = OmegaConf.create(dict_obj)
    with open(file_path, 'w') as f:
        OmegaConf.save(config=conf, f=f)

# %% ../../src/utils/config_loader.ipynb #debf2c16-b2c4-4aa1-a52e-cdd8aafb4ba3
def get_obj_from_str(string, reload=False, invalidate_cache=True):
    module, cls = string.rsplit(".", 1)
    if invalidate_cache:
        importlib.invalidate_caches()
    if reload:
        module_imp = importlib.import_module(module)
        importlib.reload(module_imp)
    return getattr(importlib.import_module(module, package=None), cls)

# %% ../../src/utils/config_loader.ipynb #ffc7fa40-81aa-42ed-ac23-8562ffdc8e4f
def instantiate_from_config(config):
    if not "target" in config: raise KeyError("Expected key `target` to instantiate.")
    if not "params" in config: print("[WARNING] Expected key `params` to instantiate.")
    return get_obj_from_str(config["target"])(**config.get("params", dict()))

# %% ../../src/utils/config_loader.ipynb #2dfa8063-d1d1-4ea6-b4ed-055e905a669f
def store_model_state_dict(state_dict, save_path):
    print(f"[INFO]: Saving model to `{save_path}`.")

    if save_path.endswith("ckpt") or save_path.endswith("pt"):
        torch.save(state_dict, save_path)

    elif save_path.endswith("safetensors"):       
        save_safetensors(state_dict, save_path)
        
    else:
        raise NotImplementedError(f"unknown filetype: {save_path}")

# %% ../../src/utils/config_loader.ipynb #06c24366-a179-492e-adc2-5d696f7485fa
def load_model_state_dict(save_path, device):
    print(f"[INFO]: Loading model from `{save_path}` onto device: {device}.")

    if save_path.endswith("ckpt") or save_path.endswith("pt"):
        state_dict = torch.load(save_path, map_location=torch.device(device).type, weights_only=True)

    elif save_path.endswith("safetensors"):       
        state_dict = load_safetensors(save_path, device=torch.device(device).type)
        
    else:
        raise NotImplementedError(f"unknown filetype: {save_path}")

    return state_dict

# %% ../../src/utils/config_loader.ipynb #ec27728c-c8a9-4adf-bdbb-cf9dd9fc4abc
def store_tensor(tensor, save_path, type="tensor"):
    print(f"[INFO]: Saving tensor to `{save_path}`.")

    if type=="numpy" and save_path.endswith("safetensors"):
        save_path = save_path.replace(".safetensors", ".pt")
    
    if save_path.endswith("ckpt") or save_path.endswith("pt") or type=="numpy":
        # serializing a string larger than 4 GiB requires pickle protocol 4 or higher; Protocol version 5 was added in Python 3.8.
        torch.save(tensor, save_path, pickle_protocol=5)   

    elif save_path.endswith("safetensors") and type=="tensor":       
        save_safetensors(tensor, save_path)
            
    else:
        raise NotImplementedError(f"unknown filetype: {save_path} or unknown type {type}")

# %% ../../src/utils/config_loader.ipynb #cc632206-1779-462d-968d-1debdbe83c91
def load_tensor(save_path, device, type="tensor"):
    print(f"[INFO]: Loading tensor from `{save_path}` onto device: {device}.")

    if type=="numpy" and save_path.endswith("safetensors"):
        save_path = save_path.replace(".safetensors", ".pt")
    
    if save_path.endswith("ckpt") or save_path.endswith("pt") or type=="numpy":
        tensor = torch.load(save_path, map_location=torch.device(device).type, weights_only=False)

    elif save_path.endswith("safetensors") and type=="tensor":            
        tensor = load_safetensors(save_path, device=torch.device(device).type)

    else:
        raise NotImplementedError(f"unknown filetype: {save_path} or unknown type {type}")

    return tensor
