import torch


def make_weights(diag_val, offdiag_val, n_units):
    """Get a connection weight matrix with "diag-offdial structure"

    e.g.
        | x, y, y |
        | y, x, y |
        | y, y, x |

    where x = diag_val, and y = offdiag_val

    Parameters
    ----------
    diag_val : float
        the value of the diag entries
    offdiag_val : float
        the value of the off-diag entries
    n_units : int
        the number of LCA nodes

    Returns
    -------
    2d array
        the weight matrix with "diag-offdial structure"

    """
    diag_mask = torch.eye(n_units)
    offdiag_mask = torch.ones((n_units, n_units)) - torch.eye(n_units)
    weight_matrix = diag_mask * diag_val + offdiag_mask * offdiag_val
    return weight_matrix.float()