im          = double(imread('noisy_input.png'));
[nx,ny,nc]   = size(im);
f           = im(:);

e                  = ones(ny, 1);
Dx_tilde           = spdiags([e -e], -1:0, ny, ny);
Dx_tilde(:, end)   = 0;
e                  = ones(nx, 1);
Dy_tilde           = spdiags([-e e], 0:1, nx, nx);
Dy_tilde(end, end) = 0;
Dx                 = kron(Dx_tilde', speye(nx));
Dy                 = kron(speye(ny), Dy_tilde);
D                  = cat(1, kron(speye(3), Dx), kron(speye(3), Dy));

lambda   = 0.030;
L       = normest(D)^2/lambda;
tau     = 1/L;
Iter    = 100;

q       = zeros(2*length(f),1);
energy       = zeros(Iter,1);

%% Iteration

for k=1:Iter

    % gradient descent
    grad        = D*(D'*q/lambda-f);
    q           = q-tau*grad;
    
    % projection
    q_proj          = reshape(q,nx*ny,2*nc);
    normq       = sqrt(sum(q_proj.^2,2));
    idx         = (normq>1);
    q_proj(idx,:)   = q_proj(idx,:)./(repmat(normq(idx),1,2*nc));
    q           = q_proj(:);

    % energy
    energy(k)        = lambda/2 * norm(D'*q/lambda-f)^2;
end


%% Plot

% solution to primal problem
u    = -1/lambda*D'*q+f;

figure;
plot(1:Iter, energy);
figure;
imshow([uint8(im), reshape(uint8(u),nx,ny,nc)]);