from typing import Optional
import torch


def file2ckpt(path: str, device: Optional[str] = None) -> dict:
    """
    Load the ckpt file into a dictionary to restart a past simulation. It is a thin wrapper around torch.load.

    Args:
        path: A string specifying the location of the ckpt file (required)
        device: A string either "cuda" or "cpu" specifying the device of the intended (optional).
            This is usefull when a model trained on a GPU machine need to be loaded into a CPU machine or viceversa.

    Examples:
        >>> # to load a ckpt of a model which was trained on GPU into a CPU machine.
        >>> ckpt = file2ckpt(path="pretrained_GPU_model.pt", device="cpu")
        >>> vae = CompositionalVae(params=ckpt.get("params"))
    """
    if device is None:
        ckpt = torch.load(path)
    elif device == 'cuda':
        ckpt = torch.load(path, map_location="cuda:0")
    elif device == 'cpu':
        ckpt = torch.load(path, map_location=torch.device('cpu'))
    else:
        raise Exception("device is not recognized")
    return ckpt