rng(42);

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

kernel = fspecial('motion', 15, 45);

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

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

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

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

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));



lambda = 0.3;

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(B).^2);

max_outer_it = 100;
energy = zeros(max_outer_it, 1);
%%
% outer proxgrad iteration
for j=1:max_outer_it
    Du = reshape(D*u, [ny*nx, 2*nc]);
    norm_Du = sqrt(sum(Du.^2, 2));
    energy(j) = 0.5 * sum((B * u - f_blurred).^2) + lambda * ...
             sum(norm_Du);
    fprintf('it=%d en=%f\n', j, energy(j));
    
    inner_lambda = (1 / (outer_tau * lambda));
    inner_f = u - outer_tau * B' * (B * 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;
    imshow(reshape(u, [ny,nx,nc]));
    
end