import importlib
import os
import sys
from typing import Callable, Dict, Union

import numpy as np
import yaml
import torch


def merge_a_into_b(a, b):
    # merge dict a into dict b. values in a will overwrite b.
    for k, v in a.items():
        if isinstance(v, dict) and k in b:
            assert isinstance(
                b[k], dict
            ), "Cannot inherit key '{}' from base!".format(k)
            merge_a_into_b(v, b[k])
        else:
            b[k] = v


def load_config(config_file):
    with open(config_file, "r") as reader:
        config = yaml.load(reader, Loader=yaml.FullLoader)
    if "inherit_from" in config:
        base_config_file = config["inherit_from"]
        base_config_file = os.path.join(
            os.path.dirname(config_file), base_config_file
        )
        assert not os.path.samefile(config_file, base_config_file), \
            "inherit from itself"
        base_config = load_config(base_config_file)
        del config["inherit_from"]
        merge_a_into_b(config, base_config)
        return base_config
    return config

def get_cls_from_str(string, reload=False):
    module_name, cls_name = string.rsplit(".", 1)
    if reload:
        module_imp = importlib.import_module(module_name)
        importlib.reload(module_imp)
    return getattr(importlib.import_module(module_name, package=None), cls_name)

def init_obj_from_dict(config, **kwargs):
    obj_args = config["args"].copy()
    obj_args.update(kwargs)
    for k in config:
        if k not in ["type", "args"] and isinstance(config[k], dict) and k not in kwargs:
            obj_args[k] = init_obj_from_dict(config[k])
    try:
        obj = get_cls_from_str(config["type"])(**obj_args)
        return obj
    except Exception as e:
        print(f"Initializing {config} failed, detailed error stack: ")
        raise e

def init_model_from_config(config, print_fn=sys.stdout.write):
    kwargs = {}
    for k in config:
        if k not in ["type", "args", "pretrained"]:
            sub_model = init_model_from_config(config[k], print_fn)
            if "pretrained" in config[k]:
                load_pretrained_model(sub_model,
                                      config[k]["pretrained"],
                                      print_fn)
            kwargs[k] = sub_model
    model = init_obj_from_dict(config, **kwargs)
    return model

def merge_load_state_dict(state_dict,
                          model: torch.nn.Module,
                          output_fn: Callable = sys.stdout.write):
    model_dict = model.state_dict()
    pretrained_dict = {}
    mismatch_keys = []
    for key, value in state_dict.items():
        if key in model_dict and model_dict[key].shape == value.shape:
            pretrained_dict[key] = value
        else:
            mismatch_keys.append(key)
    output_fn(f"Loading pre-trained model, with mismatched keys {mismatch_keys}")
    model_dict.update(pretrained_dict)
    model.load_state_dict(model_dict, strict=True)
    return pretrained_dict.keys()


def load_pretrained_model(model: torch.nn.Module,
                          pretrained: Union[str, Dict],
                          output_fn: Callable = sys.stdout.write):
    if not isinstance(pretrained, dict) and not os.path.exists(pretrained):
        output_fn(f"pretrained {pretrained} not exist!")
        return
    
    if hasattr(model, "load_pretrained"):
        model.load_pretrained(pretrained, output_fn)
        return

    if isinstance(pretrained, dict):
        state_dict = pretrained
    else:
        state_dict = torch.load(pretrained, map_location="cpu")

    if "model" in state_dict:
        state_dict = state_dict["model"]
    
    merge_load_state_dict(state_dict, model, output_fn)

def pad_sequence(data, pad_value=0):
    if isinstance(data[0], (np.ndarray, torch.Tensor)):
        data = [torch.as_tensor(arr) for arr in data]
    padded_seq = torch.nn.utils.rnn.pad_sequence(data,
                                                 batch_first=True,
                                                 padding_value=pad_value)
    length = np.array([x.shape[0] for x in data])
    return padded_seq, length