import torch


def get_model(weight_path: str = "", cuda: bool = False):
    """The model used for both SARCOS and KUKA experiments.
    
    Args:
        weight_path: Path to pre-trained weights. Returns untrained model if empty.
        cuda: Set to `True` if model was pre-trained on a GPU.

    Returns:
        PyTorch sequential model.
    """
    d_in = 21
    h = 200
    d_out = 7
    model = torch.nn.Sequential(
        torch.nn.Linear(d_in, h),
        torch.nn.ReLU(),
        torch.nn.Linear(h, h),
        torch.nn.ReLU(),
        torch.nn.Linear(h, h),
        torch.nn.ReLU(),
        torch.nn.Linear(h, h),
        torch.nn.ReLU(),
        torch.nn.Linear(h, d_out))
    if weight_path:
        # models have a Dropout layer which is missing
        # in this model, so we need to renaim the last layer in
        # order to load his trained weights.
        try:
            model.load_state_dict(torch.load(weight_path,
                                             map_location=torch.device('cuda' if cuda else 'cpu')))
        except RuntimeError:
            state_dict = torch.load(weight_path,
                                    map_location=torch.device('cuda' if cuda else 'cpu'))["state_dict"]
            state_dict["8.weight"] = state_dict.pop("9.weight")
            state_dict["8.bias"] = state_dict.pop("9.bias")
            model.load_state_dict(state_dict)
    return model