figure;
%%
im              = imread('Lenna.png');
[nx,ny,nc]      = size(im);
imd             = im2double(im);
k               = 7;
n               = nx*ny;
[~,cols,~,f] = kmeans(reshape(imd,n,nc), k);

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);
%%
% Group all x- and y- derivatives so that
% the corresponding dual variable can be reshaped easily for projection
K = [kron(speye(k), Dx); 
     kron(speye(k), Dy); 
     repmat(speye(n), 1, k)];

alpha       = 0.03;
L           = normest(K); 
sigma       = 1/L;
tau         = 1/L;
iter        = 50;

p           = zeros(2*n*k+n,1);
f           = f(:);
u           = zeros(k*n,1);
u_bar       = u;
%%
for i = 1:iter
    i
    %update p;
    p = p + sigma * K * u_bar;

    %prox the two components of p separately as they decouple
    p1 = p(1:2*k*n);
    p2 = p(2*k*n+1:end);

    %prox of sigma F^*(p1);
    p_proj = reshape(p1, n*k, 2);
    normp = sqrt(sum(p_proj.^2, 2));
    idx = (normp>alpha);
    p_proj(idx,:) =  p_proj(idx,:)./(repmat(normp(idx),1,2))*alpha;
    p1 = p_proj(:);

    % prox of sigma F^*(p2);
    p2 = p2-sigma;

    p = [p1; p2];

    %update u;
    u_old = u;
    u = u - tau*K'*p;
    u = u - tau*f;
    u(u < 0) = 0;

    %overrelax u;
    u_bar = 2 * u - u_old;
    
    [~,cmax] = max(reshape(u,n,k),[],2);
    cartoon = reshape(cols(cmax,:), nx, ny, nc);
    imshow([imd,cartoon]);
end