import torch


def get_target_net_weight_masks(target_net, weight_type='rec', device=None):
    """Return a mask for the recurrent main network weights that selects
    either recurrent or feedforward weights.

    Depending on the type of weight requested (`rec` or `ff`),
    this function returns a list of tensors akin to ``target_net.param_shapes``
    with values equal to 1 for weights that are of the requested type (
    recurrent or feedforward respectively), and 0 for the other weights.

    Args:
        target_net: The target recurrent network.
        weight_type (optional, str): The type of weight of interest. Options
            are `recurrent` or `forward`.
        device (optional): PyTorch device of returned mask.

    Return:
        (list): The corresponding mask for each tensor in the main network.
    """
    weight_masks = []
    for i, weights_shape in enumerate(target_net.param_shapes):

        weight_is_recurrent = 'info' in target_net.param_shapes_meta[i].keys() \
                        and target_net.param_shapes_meta[i]['info'] == 'hh'

        # Masks are zeros unless the current weights are of the requested type.
        mask = torch.zeros(weights_shape)
        if weight_type == 'rec' and weight_is_recurrent:
            mask = torch.ones(weights_shape)
        elif weight_type == 'ff' and not weight_is_recurrent:
            mask = torch.ones(weights_shape)
        if device is not None:
            mask = mask.to(device)
        weight_masks.append(mask)

    return weight_masks