import logging
import torch


def restore_checkpoint(model, optimizer, checkpoint_file, device):
    """
    Restores model and optimizer from a checkpoint file and returns checkpoint information.
    Has side effect of loading the state_dict for model and optimizer (i.e. modifies the instances).
    :param model: [class], torch model instance
    :param optimizer: [class], torch optimizer instance
    :param checkpoint_file: string, full file path
    :param device: [class], torch device instance
    :return: Tuple of the checkpoint values
    """
    assert checkpoint_file

    logging.info("** ** * Restore from checkpoint: %s" % checkpoint_file)
    checkpoint_state = torch.load(checkpoint_file, map_location=device)
    model.load_state_dict(checkpoint_state["model_state_dict"])
    optimizer.load_state_dict(checkpoint_state["optimizer_state_dict"])
    last_epoch = checkpoint_state["epoch"]
    global_step = checkpoint_state["global_step"]

    logging.info("  RESTORED AT epoch:%d-%s, global_step:%d" % (last_epoch, global_step))
    logging.info("** ** * Model restored! ** ** * ")

    # model.train()  # Do this in calling code for now, maybe want model.eval() there instead

    return last_epoch, global_step