from . import models
from common.base import SRBase
import torch
import numpy as np 
from pathlib import Path
import inspect, sys

AVAILABLE_MODELS = {}
for module_name in sys.modules.keys():
    if 'models.' in module_name:
        AVAILABLE_MODELS.update({
            k:v for k,v in inspect.getmembers(sys.modules[module_name], 
            lambda x: inspect.isclass(x) and SRBase in inspect.getmro(x) and (not 'base' in x.__name__.lower()) )
        })

def SaveCheckpoint(model, path):
    model_container = {
        'model': model.__class__.__name__,
        'state_dict': model.state_dict(),
        **{k:v for k,v in model.__dict__.items() if k[0] != '_' and k != 'training'}
    }
    torch.save(model_container, path)

def LoadCheckpoint(model_path):
    model_path = Path(model_path).absolute()
    if model_path.exists():
        model_container = torch.load(model_path)
        init_arg_names = list(inspect.signature(AVAILABLE_MODELS[model_container['model']]).parameters.keys())
        model = AVAILABLE_MODELS[model_container['model']](**{k:model_container[k] for k in init_arg_names})
        model.load_state_dict(model_container['state_dict'], strict=True)
        return model
    else:
        raise Exception(f"Path {model_path} does not exist.")