% WAVELET TEST  An algorithmic comparison on a two-dimensional image with 
% a wavelet basis used to turn the image into a compressible signal.
%
% * Note that Pascal Getreuer's "Wavelet Transforms in MATLAB" code is
% required for the execution of this script, and is available at
% http://www.mathworks.com/matlabcentral/fileexchange/11133 (as of April 2,
% 2009) *


%
% Coded by: Justin Ziniel, Ohio State Univ.
% Email: ziniel.1@osu.edu
% Last modified: April 2, 2009
% FBMP version 1.2 
% Copyright (c) Philip Schniter, Lee C. Potter, and Justin Ziniel, 2008
%


%% Load image, and specify wavelet basis and other required parameters

% Load the normalized cameraman image
X = load('mondrian');
X = getfield(X, 'X');

% Reduce image size to 128-by-128, and renormalize
X = imresize(X, 0.25);

X = X - min(min(X));
X = X/max(max(X));
n = length(X);

N = n^2;

% Required input parameters
wavetype = '2D db1';     	% Wavelet type
keep_frac = [0.75, 0.40, 0.25];     % # measurements / # coeffs at each level
tau = 0.02;                 % GPSR penalty term
levels = log(n)/log(2);     % Number of wavelet decomposition levels
L0 = levels - 2;            % Starting coarseness level to begin CS recovery
D_max = 2;                  % Max number of allowable repeated greedy searches

% Wavelet decomposition of input image
coeffs = wavelet(wavetype, levels, X);

% Store the L0-1 coarseness levels in different recoveries
fbmp_coeffs = zeros(size(coeffs));
fbmp_coeffs(1:2^(L0-1),1:2^(L0-1)) = coeffs(1:2^(L0-1),1:2^(L0-1));
gpsr_coeffs = zeros(size(coeffs));
gpsr_coeffs(1:2^(L0-1),1:2^(L0-1)) = coeffs(1:2^(L0-1),1:2^(L0-1));

% Initial hyperparameter choices for FBMP
p1 = 0.05;
mus = [0; 0];
sig2s = [0; .10];
sig2w = 0.01;


%% Recovery

path(path, '../functions')

for L = L0:levels
    % Grab the wavelet coefficients from the SW, NE, and SE "panes" of the
    % decomposition for the appropriate level, and vectorize
    SW = coeffs((2^(L-1)+1):(2^L), 1:(2^(L-1)));
    NE = coeffs(1:(2^(L-1)), (2^(L-1)+1):(2^L));
    SE = coeffs((2^(L-1)+1):(2^L), (2^(L-1)+1):(2^L));
    nSW = numel(SW);  nNE = numel(NE);  nSE = numel(SE);
    x = [SW(:); NE(:); SE(:)];
    
    % Generate random measurement matrix
    N = numel(x);                       % Number of unknowns
    M = ceil(keep_frac(L-L0+1)*N);      % Number of measurements
    A = randn(M,N);                     % Random normal mixing matrix
    for i = 1:N;  A(:,i) = A(:,i)/norm(A(:,i));  end;    % Unit-norm columns
    y = A*x;                            % Noiseless measurements
    

    % ---------------------------------------------------------------------
    % Recover with FBMP
    tic;
    [xmmse, xmmse_star, psy_star, nu_star, T_star, d_tot, d_max, e_tot] = ...
                fbmpr_gem_refine_fxn(y, A, p1, sig2w, sig2s, mus, D_max, 0, 20, 0.05);
    toc;
    
    % Reshape MMSE solution into "panes" and insert into FBMP's wavelet
    % decomposition matrix
    SW_fbmp = reshape(xmmse(1:nSW), 2^(L-1), 2^(L-1));
    NE_fbmp = reshape(xmmse(nSW+1:nSW+nNE), 2^(L-1), 2^(L-1));
    SE_fbmp = reshape(xmmse(nSW+nNE+1:end), 2^(L-1), 2^(L-1));
    
    fbmp_coeffs((2^(L-1)+1):(2^L), 1:(2^(L-1))) = SW_fbmp;
    fbmp_coeffs(1:(2^(L-1)), (2^(L-1)+1):(2^L)) = NE_fbmp;
    fbmp_coeffs((2^(L-1)+1):(2^L), (2^(L-1)+1):(2^L)) = SE_fbmp;
    
    
    % ---------------------------------------------------------------------
    % Recover with GPSR
    
    tic;
    [x_gpsr, x_gpsr_debias] = GPSR_Basic(y, A, tau, 'Debias', 1, ...
                               'ToleranceA', 1e-4, 'Verbose', 0);
    toc;
    % Select the recovery that minimizes MSE
    if ~isempty(x_gpsr_debias)
        [dummy, min_index] = min([norm(x - x_gpsr), norm(x - x_gpsr_debias)]);
        if min_index == 2       % Replace x_gpsr with debiased
            x_gpsr = x_gpsr_debias;
        end
    end
        
    % Reshape MMSE solution into "panes" and insert into FBMP's wavelet
    % decomposition matrix
    SW_gpsr = reshape(x_gpsr(1:nSW), 2^(L-1), 2^(L-1));
    NE_gpsr = reshape(x_gpsr(nSW+1:nSW+nNE), 2^(L-1), 2^(L-1));
    SE_gpsr = reshape(x_gpsr(nSW+nNE+1:end), 2^(L-1), 2^(L-1));
    
    gpsr_coeffs((2^(L-1)+1):(2^L), 1:(2^(L-1))) = SW_gpsr;
    gpsr_coeffs(1:(2^(L-1)), (2^(L-1)+1):(2^L)) = NE_gpsr;
    gpsr_coeffs((2^(L-1)+1):(2^L), (2^(L-1)+1):(2^L)) = SE_gpsr;
end

% Recover images from the estimates
fbmp_image = wavelet_xform(wavelet, -levels, fbmp_coeffs);
gpsr_image = wavelet_xform(wavelet, -levels, gpsr_coeffs);

% Determine NMSE (in dB) for the various recoveries
nmse_mask = ones(n,n); nmse_mask(1:2^(L0-1),1:2^(L0-1)) = zeros(2^(L0-1));
nmse_fbmp = 10*log10(norm((coeffs.*nmse_mask - fbmp_coeffs.*nmse_mask), ...
                          'fro')^2 / norm(coeffs.*nmse_mask, 'fro')^2);
nmse_gpsr = 10*log10(norm((coeffs.*nmse_mask - gpsr_coeffs.*nmse_mask), ...
                          'fro')^2 / norm(coeffs.*nmse_mask, 'fro')^2);

% Display NMSE and sparsity
disp(sprintf('FBMP NMSE: %f dB  |  FBMP sparsity: %2.2f', ...
    nmse_fbmp, 100*sum(sum((fbmp_coeffs ~= 0) .* nmse_mask))/(n^2 - 2^(2*(L0 - 1)))));
disp(sprintf('GPSR NMSE: %f dB  |  GPSR sparsity: %2.2f', ...
    nmse_gpsr, 100*sum(sum((gpsr_coeffs ~= 0) .* nmse_mask))/(n^2 - 2^(2*(L0 - 1)))));

% Plot images 
figure(1)
subplot(131)
imshow(X, []); title('(a) Original Image')
subplot(132)
imshow(fbmp_image, []); title('(b) FBMP recovery')
subplot(133)
imshow(gpsr_image, []); title('(c) GPSR recovery')
