function varargout = BayesDESIGN_SB_SWT(Y_GRAPPA,Lambda_G,cov_inds,sz_ACS,Ry,Rz,gamma,varargin)
%
% function [Y_DESIGN,info,GRAPPA_diff_normsqd,L1_norm] = BayesDESIGN_SB_SWT(Y_GRAPPA,Lambda_G,cov_inds,sz_ACS,Ry,Rz,gamma,varargin)
%
% Perform Bayesian DESIGN denoising for GRAPPA using split-Bregman
% iterations. The Split-Bregman approach solves
% [X^(k),W^(k)] = min_{X,W} 0.5*||X-X_G||_{Lambda_G}^2 + gamma*||W||_{1,2}
%                            + lambda/2*||W - SWT*IDFT*(Kbar^T*X+K^T*D) -
%                            B^(k-1)||_F^2
% B^(k) = B^(k-1) + SWT*IDFT*(Kbar^T*X^(k)+K^T*D) - W^(k)
% Pseudocode of algorithm:
%
% 1. Initialize X^(0) (missing k-space) to zero, W^(0) to
%    SWT*IDFT(Kbar^T*X+K^T*D). B^(0) = 0, choose lambda = 1. 
% 2. Generate missing k-space estimates X_G from GRAPPA recon and noise
%    covariance Lambda_G.
% 3. Get X^(k) from X^(k-1) by directly solving:
%    A*vec(X^(k)) = b, where
%    
%    A*vec(X) = Lambda_G^-1*vec(X) + lambda*vec(Kbar*IDFT'*SWT'*SWT*IDFT*Kbar^T*X)
%
%    and
%
%    b = Lambda_G^-1*vec(X_G) + lambda*vec(Kbar*IDFT'*SWT'*(W^(k-1)-B^(k-1)-SWT*IDFT*K^T*D))
%
%    Note that SWT is shift-invariant, so IDFT'*SWT'*SWT*IDFT is diagonal.
%
% 4. Perform "shrinkage" thresholding on W:
%    [W^(k)]_{n,:} = shrink([SWT*IDFT*(Kbar^T*X^(k)+K^T*D)+B^(k-1)]_{n,:},gamma/lambda),
%    where
%
%    shrink(x,gamma/lambda) = max(1-gamma/(lambda*norm(x)),0)*x
%
% 5. "Add back the noise":
%
%    B^(k) = B^(k-1) + SWT*IDFT*(Kbar^T*X^(k)+K^T*D)-W^(k)
%
% 6. Repeat 3-5 until stopping criteri(on/a) is reached (function change
%    less than delta or [W;X] change less than epsilon).
%
% Copyright 2012-10-09, Daniel Weller, University of Michigan

if isequal(Y_GRAPPA,'test'), BayesDESIGN_SB_SWT_test(); return; end

opts.preservedata = true; % uses "nullspace" method to preserve data if true
opts.lambda = 100*max(1,gamma); % Bregman penalty parameter
opts.maxiters = 20;
opts.forcemaxiters = false;
opts.delta = 1e-4; % tolerance in f
opts.epsilon = 1e-4; % tolerance in X
opts.wname = 'haar'; % wavelet name (see wfilters())
opts.wlevels = 2; % SWT levels
opts.NFirsty = []; % y-coord of first acquired
opts.NFirstz = []; % z-coord of first acquired
opts.ACSFirstx = []; % x-coord of first ACS
opts.ACSFirsty = []; % y-coord of first ACS
opts.ACSFirstz = []; % z-coord of first ACS

opts = vararg_pair(opts,varargin);

% setup dimensions and mask
[Nx,Ny,Nz,P] = size(Y_GRAPPA);
N = Nx*Ny*Nz;
if opts.preservedata
    K = false(Nx,Ny,Nz);
    if isempty(opts.NFirsty), opts.NFirsty = mod(floor(Ny/2),Ry)+1; end
    if isempty(opts.NFirstz), opts.NFirstz = mod(floor(Nz/2),Rz)+1; end
    K(:,opts.NFirsty:Ry:end,opts.NFirstz:Rz:end) = true;
    if ~isempty(sz_ACS)
        NACSx = sz_ACS(1); NACSy = sz_ACS(2); NACSz = sz_ACS(3);
        if isempty(opts.ACSFirstx), opts.ACSFirstx = floor(Nx/2) - floor(NACSx/2) + 1; end
        if isempty(opts.ACSFirsty), opts.ACSFirsty = floor(Ny/2) - floor(NACSy/2) + 1; end
        if isempty(opts.ACSFirstz), opts.ACSFirstz = floor(Nz/2) - floor(NACSz/2) + 1; end
        K(opts.ACSFirstx+(0:NACSx-1),opts.ACSFirsty+(0:NACSy-1),opts.ACSFirstz+(0:NACSz-1)) = true;
    end
    [K,Kbar] = make_mask_fatrices(K);
    M = size(K,1); % = sum(K(:))
else % no masking
end

% setup GRAPPA output and split into acquired and "missing" k-space pieces
Y_GRAPPA = reshape(Y_GRAPPA,[N,P]);
if opts.preservedata
    D = K*Y_GRAPPA;
    X_GRAPPA = Kbar*Y_GRAPPA;
else % use Y_GRAPPA directly
end

% setup DFT and DWT fatrices
[~,IDFT] = make_DFT_fatrices('im_dims',[Nx,Ny,Nz],'dftdims',1:3,'dftshift',true,'imgshift',true);
SWT = make_SWT_fatrices('im_dims',[Nx,Ny,Nz],'dwtdims',[2,3],'nlevel',opts.wlevels,'wname',opts.wname);

% get covariance info
cov_inds = cov_inds(:);
if opts.preservedata
    Xcov_inds = Kbar*cov_inds;
else % use cov_inds directly
end
Lambda_G_inv = complex(zeros(size(Lambda_G),class(Lambda_G)));
Lambda_G_invsqrt = complex(zeros(size(Lambda_G),class(Lambda_G)));
for ilambda = 1:size(Lambda_G,3)
    [eigvecs,eigvals] = eig(Lambda_G(:,:,ilambda));
    eigvals = abs(diag(eigvals));
    eigvals_max = max(eigvals);
    eigvals_invert = eigvals > eps(eigvals_max);
    eigvals(eigvals_invert) = 1./eigvals(eigvals_invert);
    eigvals(~eigvals_invert) = 0; % truncate smallest eigenvalues
    Lambda_G_inv(:,:,ilambda) = eigvecs*diag(eigvals)*eigvecs';
    Lambda_G_invsqrt(:,:,ilambda) = eigvecs*diag(sqrt(eigvals))*eigvecs';
end

% initialize X/Y, W, and B
t0 = tic();
if opts.preservedata
    X = complex(zeros([N-M,P],class(D))); % initialize "missing" k-space
    W_from_D = SWT*(IDFT*(K'*D));
    W = SWT*(IDFT*(Kbar'*X))+W_from_D;
else
    Y = complex(zeros([N,P],class(Y_GRAPPA))); % initialize full k-space
    W = SWT*(IDFT*Y);
end
B = complex(zeros(size(W),class(W)));
if opts.preservedata
    f = evalf_XWB(X,W,X_GRAPPA,Lambda_G_invsqrt,Xcov_inds,gamma);
else % f is essentially the same function in both cases
    f = evalf_XWB(Y,W,Y_GRAPPA,Lambda_G_invsqrt,cov_inds,gamma);
end
fvals = [f;zeros(opts.maxiters,1)];

% set up inversion for X update
if opts.preservedata
    Ainv_X = setup_Ainv_X(Lambda_G_inv,Xcov_inds,opts.lambda,Kbar,IDFT,SWT,[Nx,Ny,Nz,P]);
else
    Ainv_X = setup_Ainv_X(Lambda_G_inv,cov_inds,opts.lambda,Kbar,IDFT,SWT,[Nx,Ny,Nz,P]);
end

% main loop (split-Bregman)
niter = 1; converged = 0;
if opts.preservedata
    Y_DESIGN = Kbar'*X+K'*D; 
else
    Y_DESIGN = Y;
end
fmin = f; niter_min = 0;
while niter <= opts.maxiters
    % perform x update
    if opts.preservedata
%         Afun = @(in) evalA_X(in,Lambda_G_inv,Xcov_inds,opts.lambda,Kbar,IDFT,SWT);
        bvec = apply_Lambda_G_inv(X_GRAPPA,Lambda_G_inv,Xcov_inds) ...
            + opts.lambda.*(Kbar*(IDFT'*(SWT'*(W-B-W_from_D))));
        X_prev = X(:);
%         [X,CG_info_new] = pre_conj_grad(X_prev,Afun,bvec(:),CG_opts{:});
%         X = reshape(X,[N-M,P]);
        X = evalAinv_X(bvec,Ainv_X);
    else % CG is essentially the same
%         Afun = @(in) evalA_X(in,Lambda_G_inv,cov_inds,opts.lambda,1,IDFT,SWT);
        bvec = apply_Lambda_G_inv(Y_GRAPPA,Lambda_G_inv,cov_inds) ...
            + opts.lambda.*(IDFT'*(SWT'*(W-B)));
        Y_prev = Y(:);
%         [Y,CG_info_new] = pre_conj_grad(Y_prev,Afun,bvec(:),CG_opts{:});
%         Y = reshape(Y,[N,P]);
        Y = evalAinv_X(bvec,Ainv_X);
    end
    
    % perform shrinkage/thresholding on [SWT*IDFT*(Kbar^T*X^(k)+K^T*D)+B^(k-1)]
    W_prev = W(:);
    if opts.preservedata
        W_from_X = SWT*(IDFT*(Kbar'*X))+W_from_D;
        W_plus_B = W_from_X + B;
    else % Y contains all k-space
        W_from_Y = SWT*(IDFT*Y);
        W_plus_B = W_from_Y + B;
    end
    W = vector_shrinkage(W_plus_B,gamma/opts.lambda);
    
    % update B
%     B_prev = B;
    B = W_plus_B - W;
    
    % update function value
    f_prev = f;
    if opts.preservedata
        f = evalf_XWB(X,W_from_X,X_GRAPPA,Lambda_G_invsqrt,Xcov_inds,gamma);
    else % still essentially the same
        f = evalf_XWB(Y,W_from_Y,Y_GRAPPA,Lambda_G_invsqrt,cov_inds,gamma);
    end
    fvals(niter+1) = f;
    if opts.forcemaxiters || f < fmin
        if opts.preservedata
            Y_DESIGN = Kbar'*X+K'*D;
        else
            Y_DESIGN = Y;
        end
        fmin = f;
        niter_min = niter;
    end
    
    % check stopping criteria
    if opts.preservedata
        W_XY_diff = norm([W(:)-W_prev;X(:)-X_prev]);
        normWXY = norm([W_prev;X_prev]);
    else
        W_XY_diff = norm([W(:)-W_prev;Y(:)-Y_prev]);
        normWXY = norm([W_prev;Y_prev]);
    end
    f_diff = f_prev - f;
    if W_XY_diff <= opts.epsilon*normWXY
        converged = 1; % [W;X/Y] converged
    elseif f_diff <= opts.delta*f_prev
        converged = 2; % f converged
    end
    
    if ~opts.forcemaxiters && converged > 0, break; end
    
    niter = niter + 1; drawnow;
end

if niter > opts.maxiters, niter = opts.maxiters; end

if opts.preservedata
    GRAPPA_diff_normsqd = 0.5*sum(sum(abs(apply_Lambda_G_inv(Kbar*Y_DESIGN-X_GRAPPA,Lambda_G_invsqrt,Xcov_inds)).^2,2),1);
else % still essentially the same
    GRAPPA_diff_normsqd = 0.5*sum(sum(abs(apply_Lambda_G_inv(Y_DESIGN-Y_GRAPPA,Lambda_G_invsqrt,cov_inds)).^2,2),1);
end
L1_norm = sum(sqrt(sum(abs(SWT*(IDFT*Y_DESIGN)).^2,2)),1);

Y_DESIGN = reshape(Y_DESIGN,[Nx,Ny,Nz,P]);
ttotal = toc(t0);

info.niters = niter;
info.niter_min = niter_min;
info.ttotal = ttotal;
info.converged = converged;
info.fvals = fvals(1:niter+1);

% l1-l2 term for W

if nargout >= 1, varargout{1} = Y_DESIGN; end
if nargout >= 2, varargout{2} = info; end
if nargout >= 3, varargout{3} = GRAPPA_diff_normsqd; end
if nargout >= 4, varargout{4} = L1_norm; end

end

%
% function X = apply_Lambda_G_inv(X,Lambda_G_inv,Xcov_inds)
%
% Apply appropriate Lambda_G^-1 to each row of 2-D matrix X according to index
% specified in Xcov_inds. Output is same size as X. Xcov_inds should be
% column vector same size as a column of X.
%
function X = apply_Lambda_G_inv(X,Lambda_G_inv,Xcov_inds)

N = size(X,1);
if length(Xcov_inds) ~= N, error('Xcov_inds wrong size!'); end
for ir = 1:size(Lambda_G_inv,3)
    do_r = find(Xcov_inds == ir);
    if isempty(do_r), continue; end
    X_proc = X(do_r,:).';
    X_proc = Lambda_G_inv(:,:,ir)*X_proc;
    X(do_r,:) = X_proc.';
end

end

function f = evalf_XWB(X,W_from_X,X_GRAPPA,Lambda_G_invsqrt,Xcov_inds,gamma)

% least-squares term for X
f = 0.5*sum(sum(abs(apply_Lambda_G_inv(X-X_GRAPPA,Lambda_G_invsqrt,Xcov_inds)).^2,2),1);

% l1-l2 term for W
f = f + gamma*sum(sqrt(sum(abs(W_from_X).^2,2)),1);

end

%
% function out = evalA_X(in,Lambda_G_inv,Xcov_inds,lambda,Kbar,IDFT,SWT)
%
% Function evaluates A*in for CG iteration to update X.
%
function out = evalA_X(in,Lambda_G_inv,Xcov_inds,lambda,Kbar,IDFT,SWT)

N_minus_M = length(Xcov_inds);
P = size(Lambda_G_inv,1);
in = reshape(in,[N_minus_M,P]); % CG provides vectorized input
out = apply_Lambda_G_inv(in,Lambda_G_inv,Xcov_inds);
in = Kbar'*in;
for ip = 1:size(in,2);
    in(:,ip) = SWT_idftgramvec(SWT,IDFT,in(:,ip));
end
out = out + lambda.*(Kbar*in);
% out = out + lambda.*(Kbar*(IDFT'*(SWT'*(SWT*(IDFT*(Kbar'*in))))));
out = out(:); % CG expects vectorized output

drawnow;

end

% 
% function out = evalAinv_X(in,Ainv_X)
%
% Function evaluates (Lambda_G^{-1}+Kbar*IDFT'*SWT'*SWT*IDFT*Kbar')\in for updating X.
%
function out = evalAinv_X(in,Ainv_X)

out = zeros(size(in),class(in));
for ip = 1:size(Ainv_X,3)
    out(:,ip) = sum(Ainv_X(:,:,ip).*in,2);
end

end

% 
% function Ainv_X = setup_Ainv_X(Lambda_G_inv,Xcov_inds,lambda,Kbar,IDFT,SWT,sz_full)
%
% Function sets up (Lambda_G^-1+Kbar*IDFT'*SWT'*SWT*IDFT*Kbar')^-1.
%
function Ainv_X = setup_Ainv_X(Lambda_G_inv,Xcov_inds,lambda,Kbar,IDFT,SWT,sz_full)

% assemble IDFT'*TV'*TV*IDFT
SWTgram_fft = ones(sz_full(1:3));
SWTgram_fft = SWT_idftgramvec(SWT,IDFT,SWTgram_fft(:));

% subsample according to Kbar
SWTgram_fft = lambda.*(Kbar*SWTgram_fft);

% now do update for each k-space point in X:
N = size(SWTgram_fft,1);
if length(Xcov_inds) ~= N, error('Xcov_inds wrong size!'); end

Ainv_X = zeros([sz_full(4),sz_full(4),N]);

% put Lambda_G_inv(:,:,ir) in the right places
for ir = 1:size(Lambda_G_inv,3)
    do_r = find(Xcov_inds == ir);
    if isempty(do_r), continue; end
    Ainv_X(:,:,do_r) = repmat(Lambda_G_inv(:,:,ir),[1,1,length(do_r)]);
end

% add the FFT of the TV gram matrix
for ip = 1:sz_full(4)
    Ainv_X(ip,ip,:) = Ainv_X(ip,ip,:) + shiftdim(SWTgram_fft,-2);
end

% invert all the matrices
Ainv_X = arrayfun(@(x) inv(Ainv_X(:,:,x)), 1:N, 'UniformOutput',false);
Ainv_X = cat(3,Ainv_X{:});
Ainv_X = permute(Ainv_X,[3,2,1]);

end

%
% function X = vector_shrinkage(X,thresh)
%
% Perform vector soft-thresholding on X, which is solution to separable
% l1-l2 norm + least-squares term.
%
function X = vector_shrinkage(X,thresh)

P = size(X,2);
Xmags = sqrt(sum(abs(X).^2,2));
Xmags = 1-thresh./Xmags;
Xmags(Xmags < 0) = 0;
X = repmat(Xmags,[1,P]).*X;

end

%
% Test function for BayesDESIGN_SB_SWT()
%
function BayesDESIGN_SB_SWT_test()

global Y_GRAPPA GRAPPA_MSE Y_DESIGN BayesDESIGN_info DESIGN_MSE; 

rng('default');

% load one slice of 12-channel brain data
load('/y/wellerds/MRI_data/BrainWeb/T1w_sim12ch_s81.mat');

preservedata = true;
im_ref = double(im_ref); smap = double(smap); smap_norm = double(smap_norm); noisecov = double(noisecov); %#ok<NODEF>
if usejava('awt')
    figure; imshow(squeeze(im_ref),[0,1]);
    title('Reference Image'); drawnow;
%     set(gcf,'Name','Reference Image'); drawnow;
end

% make multi-channel k-space
[Nx,Ny,Nz,NCha] = size(smap);
ref = im_ref;
ref = fft3c(ref(:,:,:,ones(1,NCha)).*smap);
clear smap;

% create noise
noisevar = 0.25; % in multi-channel k-space
Lambda = noisevar.*noisecov;
noise = sqrtm(Lambda./2)*complex(randn([NCha,Nx*Ny*Nz],class(im_ref)),randn([NCha,Nx*Ny*Nz],class(im_ref)));
noise = reshape(noise.',[Nx,Ny,Nz,NCha]);
data_full = ref + noise;
clear noise;
im_full = sos_combine(ifft3c(data_full),noisecov,smap_norm);
im_diff = im_full-im_ref;
if usejava('awt')
    figure; imshow(squeeze(abs(im_diff)),[]); colorbar(); xlabel(sprintf('SER = %g dB',20*log10(norm(im_ref(:))/norm(im_diff(:)))));
    title('Additive Noise on Full Magnitude Image'); drawnow;
%     set(gcf,'Name','Additive Noise on Full Magnitude Image'); drawnow;
end

% sample k-space
Ry = 2;
Rz = 2;
NACSx = Nx; NACSy = 24; NACSz = 24;
NFirsty = mod(floor(Ny/2),Ry)+1;
NFirstz = mod(floor(Nz/2),Rz)+1;
ACSFirstx = floor(Nx/2) + 1 - floor(NACSx/2);
ACSFirsty = floor(Ny/2) + 1 - floor(NACSy/2);
ACSFirstz = floor(Nz/2) + 1 - floor(NACSz/2);
K = false(Nx,Ny,Nz);
data_redu = data_full(:,NFirsty:Ry:end,NFirstz:Rz:end,:);
K(:,NFirsty:Ry:end,NFirstz:Rz:end) = true;
ACS = data_full(ACSFirstx:ACSFirstx+NACSx-1,ACSFirsty:ACSFirsty+NACSy-1,ACSFirstz:ACSFirstz+NACSz-1,:);
K(ACSFirstx:ACSFirstx+NACSx-1,ACSFirsty:ACSFirsty+NACSy-1,ACSFirstz:ACSFirstz+NACSz-1) = true;

% perform GRAPPA calibration and reconstruction
Bx = 1; By = 4; Bz = 4;
ttotal = tic();
[Kernel,Lambda_G] = GRAPPA_construct_kernel(ACS,Ry,Rz,Bx,By,Bz,Lambda);
[Y_GRAPPA,cov_inds] = GRAPPA_recon([Nx,Ny,Nz,NCha],data_redu,Kernel,ACS,NFirsty,NFirstz,ACSFirstx,ACSFirsty,ACSFirstz);
ttotal = toc(ttotal);
Y_DIFF = Y_GRAPPA - ref;
if preservedata, Y_DIFF = Y_DIFF(~K(:,:,:,ones(1,NCha))); end
GRAPPA_MSE = norm(Y_DIFF(:))^2;
fprintf(1,'GRAPPA: MSE = %g (%g s).\n',GRAPPA_MSE,ttotal); drawnow;

% combine coils and display GRAPPA output
im_GRAPPA = sos_combine(ifft3c(Y_GRAPPA),noisecov,smap_norm);
im_diff = im_GRAPPA-im_ref;

if usejava('awt')
    figure; imshow(squeeze(im_GRAPPA),[0,1]); xlabel(sprintf('SER = %g dB',20*log10(norm(im_ref(:))/norm(im_diff(:)))));
    title('GRAPPA-Reconstructed Image'); drawnow;
%     set(gcf,'Name','GRAPPA-Reconstructed Image'); drawnow;

    figure; imshow(squeeze(abs(im_diff)),[]); colorbar(); xlabel(sprintf('SER = %g dB',20*log10(norm(im_ref(:))/norm(im_diff(:)))));
    title('GRAPPA-Reconstructed Magnitude Noise Image'); drawnow;
%     set(gcf,'Name','GRAPPA-Reconstructed Magnitude Noise Image'); drawnow;
end

% perform BayesDESIGN_SB
gammas = logspace(1,4,7);
Y_DESIGN = cell(size(gammas));
BayesDESIGN_info = cell(size(gammas));
DESIGN_MSE = zeros(size(gammas));
for ig = 1:length(gammas)
    [Y_DESIGN{ig},BayesDESIGN_info{ig}] = BayesDESIGN_SB_SWT(Y_GRAPPA,Lambda_G,cov_inds,size(ACS),Ry,Rz,gammas(ig),'preservedata',preservedata,'lambda',1e4,'maxiters',20,'forcemaxiters',true,'wname','bior4.4','wlevels',4,'NFirsty',NFirsty,'NFirstz',NFirstz,'ACSFirstx',ACSFirstx,'ACSFirsty',ACSFirsty,'ACSFirstz',ACSFirstz);
    Y_DIFF = Y_DESIGN{ig} - ref;
    if preservedata, Y_DIFF = Y_DIFF(~K(:,:,:,ones(1,NCha))); end
    DESIGN_MSE(ig) = norm(Y_DIFF(:))^2;
    fprintf(1,'BayesDESIGN (gamma = %g): MSE = %g; f = %g -> %g after %d iterations (%g s).\n',gammas(ig),DESIGN_MSE(ig),BayesDESIGN_info{ig}.fvals(1),BayesDESIGN_info{ig}.fvals(BayesDESIGN_info{ig}.niter_min+1),BayesDESIGN_info{ig}.niter_min,BayesDESIGN_info{ig}.ttotal); drawnow;

    % combine coils and display DESIGN output
    im_DESIGN = sos_combine(ifft3c(Y_DESIGN{ig}),noisecov,smap_norm);
    im_diff = im_DESIGN-im_ref;
    
    if usejava('awt')
        figure; imshow(squeeze(im_DESIGN),[0,1]); xlabel(sprintf('\\gamma = %g: SER = %g dB',gammas(ig),20*log10(norm(im_ref(:))/norm(im_diff(:)))));
        title('DESIGN-Reconstructed Image'); drawnow;
%         set(gcf,'Name','DESIGN-Reconstructed Image'); drawnow;

        figure; imshow(squeeze(abs(im_diff)),[]); colorbar(); xlabel(sprintf('\\gamma = %g: SER = %g dB',gammas(ig),20*log10(norm(im_ref(:))/norm(im_diff(:)))));
        title('DESIGN-Reconstructed Magnitude Noise Image'); drawnow;
%         set(gcf,'Name','DESIGN-Reconstructed Magnitude Noise Image'); drawnow;
    end
end

end
