import torch


def random_task(dim, use_gpu=False):
    """Create a random linear system of fixed dimension.

    For better reproducibility, the seeds are set so zero.

    Parameters:
    -----------
    dim (int): Dimension of the system
    use_gpu (bool): Store quantities on GPU, else CPU

    Returns:
    --------
    A (torch.Tensor): Random regularized PSD matrix describing the system
    b (torch.Tensor): Right hand side of the linear system
    """
    device = torch.device("cuda:0" if use_gpu else "cpu")
    # random PSD matrix
    A = torch.rand(dim, dim, device=device)
    A = A.t().matmul(A) / torch.norm(A) ** 2
    b = torch.rand(dim, device=device)
    # regularize
    diag = torch.diagflat(torch.ones_like(b))
    A = A + diag
    return A, b