im = imread('fish_saltpepper.png');
[ny,nx,nc] = size(im);
f = im2double(im);

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

%% Parameters 1
alpha = 1;
L = normest(D);

tau = 1/L; 
sigma = 1/L; 

iter = 120;

n = ny*nx*nc;
p = zeros(2*n,1);
f = f(:); 
u = f;
u_bar = u;

E_pdhg = zeros(iter, 1);

%% Iteration
for k = 1:iter

    TV = sum(sqrt(sum(reshape(D*u,ny*nx,2*nc).^2, 2)));
    en = sum(abs(u-f)) + alpha*TV
    E_pdhg(k, 1) = en;
    
    % Update p 
    p = p+sigma*D*u_bar;
    
    p_proj       = reshape(p,ny*nx,2*nc);
    normp        = sqrt(sum(p_proj.^2,2));
    idx          = (normp>alpha);
    p_proj(idx,:)= alpha*p_proj(idx,:)./(repmat(normp(idx),1,2*nc));
    p            = p_proj(:);
    
    
    u_old = u;
    % Update u
    u = u-tau*D'*p;
    
    % prox u
    v = u-f;
    u = zeros(n,1);
    u(v<-tau) = v(v<-tau)+tau;
    u(v>tau) = v(v>tau)-tau;
    u = u+f;
    %imshow(uint8(reshape(u, ny, nx, nc)*255))
    % Update u_bar
    u_bar = 2*u-u_old;
end
%%
figure; 
plot(1:iter, E_pdhg, 'red');