rng(42);

im = imresize(imread('flowers.png'), 1);
[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(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);
D                  = cat(1, kron(speye(3), Dx), kron(speye(3), Dy));

lambda = 0.03;

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;
        
        qproj = reshape(q, [ny*nx, nc*2]);
        norm_qproj = sqrt(sum(qproj.^2,2));
        qproj(norm_qproj > 1, :, :) = qproj(norm_qproj > 1, :, :) ./ ...
            repmat(norm_qproj(norm_qproj>1), 1, nc*2); 
    
        q = qproj(:);
    end
    u = inner_f - (D' * q) / inner_lambda;
    imshow(reshape(u, [ny,nx,nc]));
    
end