% Reconstruct a walnut slice from real X-ray data using Besov B111
% regularization implemented via Daubechies 2 wavelets.
%
% Samuli Siltanen September 2016

% Set path
path(path,'/Users/samulisiltanen/___tutkimus/matlab')

% Gamma correction for saved intermediate images
gammacorr= .5;

% Number of angles used
Nangvec = [120 30 20 15 12 10 6 4];
Nangvec = [20];

for aaa = 1:length(Nangvec)
    Nang = Nangvec(aaa);
    
    % Load the measurement matrix and the sinogram from file.
    % Discretization of the unknown is NxN.
    % Number of projection directions is Nang.
    disp(['Loading precomputed data with Nang=', num2str(Nang)])
    res  = 4; % This can be 1 or 2 or 4
    N    = res*82;
    eval(['load ../data/Data', num2str(N), '_', num2str(Nang), ' A m normA'])
    A = A/normA; % Normalize measurement matrix so that its norm is one
    m = m/normA; % Normalize measurement as well
    disp('Data loaded') 
    
    % Set wavelet transform depth
    D = log2(res*2);
    
    % Number of iterations
    Niter = 3000;
    
    % Threshold mu used in the transform (kind of regularization parameter)
    mu = .00006;
    
    % Convolution kernels for the Daubechies 2 wavelet
    h = [1+sqrt(3) 3+sqrt(3) 3-sqrt(3) 1-sqrt(3)]/(4*sqrt(2));
    g = [h(4)      -h(3)     h(2)      -h(1)];
    
    % Initialize iteration
    recn       = zeros(N);
    normvec = repmat(NaN,1,Niter);
    dfitvec = repmat(NaN,1,Niter);
    maxivec = repmat(NaN,1,Niter);
    sparvec = repmat(NaN,1,Niter);
    diffvec = repmat(NaN,1,Niter);
    imcounter  = 0;
    
    % Perform iteration
    for nnn = 1:Niter
        % Record previous iterate
        previous = recn;
               
        % Forward projection of the current iterate
        Arecn = A*recn(:);
        
        % Compare to measured data
        datafit = m(:) - Arecn;
        dfitvec(nnn) = norm(datafit)/norm(m(:));
        
        % Back-projection of the residual
        tmp = A.'*datafit;
        tmp = reshape(tmp,N,N);
        
        % DDD iteration f{n} = S_mu (f{n-1} + A^T(m - A f{n-1}))
        [recn,nonzero,tr,tmp] = Smu_wavelet_oper(h,g,recn+tmp,D,mu);
        sparvec(nnn) = nonzero;
        
        % Update norm vector (Besov space B111 norm)
        normvec(nnn) = sum(abs(tr(:)));
        
        % Projection to non-negative "quadrant"
        recn = max(0,recn);
        
        % Update sup norm vector
        maxivec(nnn) = max(abs(recn(:)));
        
        % Record change in iterate
        diffvec(nnn) = (norm(recn(:)-previous(:)))/norm(previous(:));

        
        % Plot the results
        if (mod(nnn,30)==1)|(nnn==Niter)
            figure(1)
            clf
            plotim = [scale01(recn),scale01(log(1+abs(tr)))];
            plotim = plotim.^(.5); % gamma correction for image shades
            imagesc(plotim)
            axis equal
            axis off
            colormap gray
            title(['Angles ', num2str(Nang), ', iteration number ', num2str(nnn)],'fontsize',16)
            
            % Monitor crucial parameters of the process
            figure(2)
            clf
            subplot(5,1,1)
            plot(1:Niter,normvec)
            xlim([1 Niter])
            title('Sum of absolute values of wavelet coefficients')
            subplot(5,1,2)
            plot(1:Niter,sparvec)
            xlim([1 Niter])
            title('Ratio of nonzero coefficients')
            subplot(5,1,3)
            semilogy(1:Niter,diffvec)
            xlim([1 Niter])
            title('Relative square norm change in iterate')
            subplot(5,1,4)
            plot(1:Niter,maxivec)
            xlim([1 Niter])
            title('Sup norm of iterate')
            subplot(5,1,5)
            semilogy(1:Niter,dfitvec)
            xlim([1 Niter])
            title('Data fit')
            drawnow
            
            
            % Save results to file
            imcounter = imcounter+1;
            imname = ['images/recon_DB2_', num2str(N), '_', num2str(Nang), '_',num2str(imcounter,'%04d'), '.png'];
            imwrite(uint8(255*scale01(recn).^gammacorr),imname,'png');
            
        end
        disp([nnn Niter])
    end
    
    % Save results to file
    eval(['save ../data/recon_DB2_', num2str(N), '_', num2str(Nang), ' recn Niter mu D sparvec normvec diffvec maxivec'])
    
end