%% Generate a random low-rank matrix
% A is the true matrix.
A = randn(300,10)*randn(10,100);

% randomly remove 20% entries of A.
F = A.*(rand(size(A))<0.8);

% find the indexes of known elements.
idx = find(F~=0);

% add some noise on the known elements.
F(idx) = F(idx) + normrnd(0,0.1,size(idx));

%% Using proximal gradient to recover A.
% initialization
alpha = 0.3;
X = zeros(size(F));
max_it = 800;
gap_tol = 1e-12;
tol = 1e-6;
step_size = 1;

primal_eng = zeros(1,max_it);
dual_eng = zeros(1,max_it);

for i = 1:max_it
    % TODO: apply proximal gradient on X
    X(idx) = X(idx)-alpha*step_size*(X(idx)-F(idx));
    [U,Sigma,V] = svd(X);
    X = U*max(0,Sigma-step_size)*V';
    
    % TODO: compute the primal energy
    eigs = svd(X);
    primal_eng(1,i) = alpha/2*sum((X(idx)-F(idx)).^2) + sum(eigs);
    
    % TODO: compute the dual variable from optimality condition
    Y = zeros(size(X));
    Y(idx) = -alpha*(X(idx)-F(idx));
    
    % TODO: compute the dual energy from dual variable
    [Uy, Sigmay, Vy] = svd(Y);
    if(Sigmay(1,1)-1>tol)
        dual_eng(1,i) = -Inf;
    else
        dual_eng(1,i) = -(F(idx))'*Y(idx) + 1/2/alpha*(Y(idx))'*(Y(idx));
    end
    
    fprintf('%d-th iteration, primal-dual gap:%d\n', i, abs(dual_eng(1,i)+primal_eng(1,i)))
    
    % If the primal-dual gap is smaller than the stop_criterion, we find
    % optimal X.
    if(abs(dual_eng(1,i)+primal_eng(1,i))<gap_tol)
        break;
    end
end