%%
rng(42);

im = imresize(imread('flowers.png'), 0.5);
[ny, nx, nc] = size(im);
f = double(im(:)) / 255.; 

kernel = fspecial('motion', 20, 50);

K = convmtx2(kernel, ny, nx);
K = kron(speye(nc), K);

kx = size(kernel, 2);
ky = size(kernel, 1);

nx2 = nx + kx - 1;
ny2 = ny + ky - 1;

f_blurred = K * f;
f_blurred = f_blurred + 0.05 * randn(ny2*nx2*nc,1);

imshow(reshape(f_blurred, [ny2,nx2,nc]));

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));
%%
alpha = 0.01;
lambda = 2;

u = zeros(nx*ny*nc, 1);
p = zeros(ny*nx*nc*2, 1);
z = zeros(ny*nx*nc*2,1);
iter = 150;

A = K'*K + lambda*(D'*D);
L = ichol(A);

energy_admm = zeros(iter, 1);
%%
tic
for i = 1:iter
    %update p;
    p = D*u + z/lambda;

    %prox of via moreau;
    
    %rescale with 1/stepsize
    %proj on 2,\infty unit ball
    pproj = reshape(p*(lambda/alpha), [ny*nx, nc*2]);
    norm_pproj = sqrt(sum(pproj.^2, 2));
    idx = (norm_pproj > 1);
    pproj(idx, :, :) = pproj(idx, :, :) ./ ...
        repmat(norm_pproj(idx), 1, nc*2); 


    %recover prox
    p = p - (alpha/lambda)*pproj(:);
    
    
    %u = (K'*K + D'*D) \ (K'*f_blurred + D'*(p-z));
    b = K'*f_blurred + D'*(lambda*p-z);
    [u,fl1,rr1,it1,rv1] = pcg(A,b,1e-3,30,L,L', u);
    
    
    z = z + lambda * (D*u - p);
    
    Du = reshape(D*u, [ny*nx, 2*nc]);
    norm_Du = sqrt(sum(Du.^2, 2));
    
    energy_admm(i) = 0.5 * sum((K * u - f_blurred).^2) + alpha * ...
             sum(norm_Du);
    
    fprintf('it=%d en=%f feas=%f\n', i, energy_admm(i), norm(D*u-p));
    imshow(reshape(u, [ny,nx,nc]));
end
toc
%%
max_inner_it = 15;

u = zeros(ny*nx*nc, 1);
q = zeros(ny*nx*nc*2, 1);
tau = 1/(normest(D).^2);
outer_tau = 1/(normest(K).^2);

energy_proxgrad = zeros(iter, 1);
%%
% outer proxgrad iteration
tic
for i=1:iter

    
    inner_lambda = (1 / (outer_tau * alpha));
    inner_f = u - outer_tau * K' * (K * u - f_blurred);
    % inner iteration, solve TV problem
    for k=1:max_inner_it   
        grad = D * (D' * q - inner_lambda * inner_f);
        q = q - tau * grad;
        
        q_proj = reshape(q,ny*nx,6);
        normq = sqrt(sum(q_proj.^2,2));
        idx = (normq>1);
        q_proj(idx,:) = q_proj(idx,:)./(repmat(normq(idx),1,2*3));
        q = q_proj(:);
    end
    u = inner_f - (D' * q) / inner_lambda;
    
    Du = reshape(D*u, [ny*nx, 2*nc]);
    norm_Du = sqrt(sum(Du.^2, 2));
    energy_proxgrad(i) = 0.5 * sum((K * u - f_blurred).^2) + alpha * ...
             sum(norm_Du);
    fprintf('it=%d en=%f\n', i, energy_proxgrad(i));
    imshow(reshape(u, [ny,nx,nc]));
end
toc
%%
figure;
hold on;
plot(1:iter, energy_admm);
plot(1:iter, energy_proxgrad);