f = imread('fish_saltpepper.png');
ny = size(f, 1);
nx = size(f, 2);
nc = size(f, 3);
f = im2double(f(:));
u = f;

% Construct finite difference gradient operator
e = ones(nx, 1);
A = spdiags([-1 * e, e], [0, -1], nx, nx);

e = ones(ny, 1);
B = spdiags([-1 * e, e], [0, 1], ny, ny);
Dx = kron(A', speye(ny));
Dy = kron(speye(nx), B);
D = cat(1, kron(eye(3), Dx), kron(eye(3), Dy));

u = f;
lambda = 1;

max_iter = 10000;
energy = zeros(max_iter, 1);
%%
for k=1:max_iter
    
    Du = reshape(D*u, [ny*nx, nc, 2]);
    norm_Du = sqrt(sum(sum(Du.^2,2), 3));

    energy(k) = lambda * sum(abs(u-f)) + sum(norm_Du);    
    
    Du_pos = norm_Du>0;
    
    g = zeros(ny*nx, nc, 2);
    g(Du_pos,:,:) = Du(Du_pos,:,:) ./ repmat(norm_Du(Du_pos), 1, nc, 2);
    
    grad = lambda * sign(u - f) + D' * g(:);
    
    alpha = 2/k;
    u = u - alpha * grad;
    
    fprintf('It=%d En=%f\n', k, energy(k));
end
%%
imshow([reshape(f, [ny,nx,nc]) reshape(u, [ny,nx,nc])]);