def saga(A, y, gradf, prox, maxiter, gamma):
    """Minimal and straightforward implementation of SAGA."""
    n, d = A.shape  # number of observations x dimensionality
    x = np.zeros(d)
    grads = np.array([gradf(A, x, y, i) for i in range(n)])
    avg = grads.mean(axis=0)
    X = np.empty((maxiter, d))

    for epoch in range(maxiter):
        #indices = np.random.permutation(n)
        indices = np.random.randint(0, n, n)

        for i in indices:
            grad = gradf(A, x, y, i)
            x = x - gamma * (grad - grads[i,:] + avg)
            x = prox(x, gamma)
            avg += (grad - grads[i,:]) / n
            grads[i,:] = grad

        X[epoch,:] = x

    return X


gradf = lambda A, x, y, i: A[i,:].T.dot(A[i,:].dot(x) - y[i])
prox_no = lambda x, gamma: x