import torch


def create_sparse_input(dictionary, K=1, num_samples=1):
    """
    Create sparse data given a dictionary and sparsity value.
    :param dictionary: Tensor. Dictionary to base the data on.
    :param K: Int. Sparsity value, use this number of atoms to create the data.
    :param num_samples: Number of samples to create.
    :return: Tensor. Created data
    """
    atoms = torch.randint(dictionary.size()[1], (num_samples, K))
    coefs = torch.randn((K, num_samples, 1))

    X = []
    for sample_ind in range(num_samples):
        input = dictionary[:, atoms[sample_ind,:]].mm(coefs[:, sample_ind])

        X.append(input)

    X = torch.stack(X, 1).squeeze(-1)

    return X, atoms, coefs