import torch


def load_model(model, model_path):
    """
    Load model from saved weights.
    """
    # import ipdb
    # ipdb.set_trace()
    if '.pt' in model_path:
        if hasattr(model, "module"):
            model.module.load_state_dict(torch.load(model_path, map_location="cpu")['module'], strict=False)
        else:
            model.load_state_dict(torch.load(model_path, map_location="cpu")['module'], strict=False)
    else:
        if hasattr(model, "module"):
            model.module.load_state_dict(torch.load(model_path, map_location="cpu"), strict=False)
        else:
            model.load_state_dict(torch.load(model_path, map_location="cpu"), strict=False)
    return model

def _load_state_dict_into_model(model_to_load, model_path, start_prefix="", lora_pretrained_model_path=None):
    # Convert old format to new format if needed from a PyTorch state_dict

    # copy state_dict so _load_from_state_dict can modify it
    state_dict = torch.load(model_path, map_location="cpu")
    metadata = getattr(state_dict, "_metadata", None)
    state_dict = state_dict.copy()
    if metadata is not None:
        state_dict._metadata = metadata

    if lora_pretrained_model_path is not None:
        lora_state_dict = torch.load(lora_pretrained_model_path, map_location="cpu")
        lora_metadata = getattr(lora_state_dict, "_metadata", None)
        lora_state_dict = lora_state_dict.copy()
        if lora_metadata is not None:
            lora_state_dict._metadata = lora_metadata

    error_msgs = []

    # PyTorch's `_load_from_state_dict` does not copy parameters in a module's descendants
    # so we need to apply the function recursively.
    def load(module, state_dict, prefix=""):
        local_metadata = {} if metadata is None else metadata.get(prefix[:-1], {})
        args = (state_dict, prefix, local_metadata, True, [], [], error_msgs)
        # Parameters of module and children will start with prefix. We can exit early if there are none in this
        # state_dict
        if len([key for key in state_dict if key.startswith(prefix)]) > 0:
            import deepspeed
            # In sharded models, each shard has only part of the full state_dict, so only gather
            # parameters that are in the current state_dict.
            named_parameters = dict(module.named_parameters(prefix=prefix[:-1], recurse=False))
            params_to_gather = [named_parameters[k] for k in state_dict.keys() if k in named_parameters]
            if len(params_to_gather) > 0:
                # because zero3 puts placeholders in model params, this context
                # manager gathers (unpartitions) the params of the current layer, then loads from
                # the state dict and then re-partitions them again
                with deepspeed.zero.GatheredParameters(params_to_gather, modifier_rank=0):
                    if torch.distributed.get_rank() == 0:
                        module._load_from_state_dict(*args)

        for name, child in module._modules.items():
            if child is not None:
                load(child, state_dict, prefix + name + ".")

    load(model_to_load, state_dict, prefix=start_prefix)
    if lora_pretrained_model_path is not None:
        load(model_to_load, lora_state_dict, prefix=start_prefix)
        del lora_state_dict
    # Delete `state_dict` so it could be collected by GC earlier. Note that `state_dict` is a copy of the argument, so
    # it's safe to delete it.
    del state_dict

    return model_to_load