function [WSURE,f_Y] = BayesDESIGN_SB_WSURE3(Y_GRAPPA,cov_inds,b,Y_GRAPPA_b,Kernel,Lambda_G,sz_ACS,gamma,epsilon,varargin)
% function [WSURE,f_Y] = BayesDESIGN_SB_WSURE3(Y_GRAPPA,cov_inds,b,Y_GRAPPA_b,Kernel,Lambda_G,sz_ACS,gamma,epsilon,varargin)
%
% Perform WSURE estimation for the sparsity tuning parameter for Bayesian
% DESIGN implemented with Split-Bregman, as a function of gamma.
% Internally, the method performs two iterations of GRAPPA +
% BayesDESIGN_SB, one with just the data, one with the data + complex-valued
% Bernoulli noise. GRAPPA is also performed on that Bernoulli noise (multiplied by the noise
% covariance matrix). This is a modification of BayesDESIGN_SB_WSURE().
%
% Modified to call either BayesDESIGN_SB() for DWT-based DESIGN or
% BayesDESIGN_SB_TV() for TV-based DESIGN, or BayesDESIGN_SB_SWT() for
% SWT-based DESIGN.
%
% Inputs:
%   Y_GRAPPA - output of GRAPPA
%   cov_inds - index into Lambda_G matrix for each output point in Y_GRAPPA
%   b - complex noise to add
%   Y_GRAPPA_b - output of GRAPPA for added complex noise
%   Kernel - GRAPPA kernel calibrated for this reconstruction
%   Lambda_G - coil and GRAPPA-amplified noise covariance matrices
%   sz_ACS - the size of the ACS lines (or empty, if not to be included in the GRAPPA output)
%   gamma - the tuning parameter factor for BayesDESIGN_SB/BayesDESIGN_SB_TV
%   epsilon - the scaling parameter for the Bernoulli noise
%
% Options: varargin can contain key/value pairs for the following
%   NFirsty,NFirstz - the initial index of the undersampled data in the y, z directions
%   ACSFirstx,ACSFirsty,ACSFirstz - the initial index of the ACS region in the x, y, and z directions
%   preservedata - true: will use nullspace BayesDESIGN; false: will denoise acquired data as well
%   BayesDESIGN_opts - cell array of key/value pairs to pass to BayesDESIGN
%   f_Y - if provided, output of a previous call of BayesDESIGN for reducing computation
%   regularizer - 'dwt', 'swt', or 'tv'
%
% Outputs:
%   WSURE - the estimate of the WSURE for this value of gamma
%   f_Y - the denoised full k-space (no additional noise)
%   Y_GRAPPA - the GRAPPA-reconstructed full k-space (no additional noise)
%
% Copyright 2012-10-16, Daniel Weller, University of Michigan

opts.NFirsty = [];
opts.NFirstz = [];
opts.ACSFirstx = [];
opts.ACSFirsty = [];
opts.ACSFirstz = [];
opts.preservedata = true;
opts.BayesDESIGN_opts = {};
opts.f_Y = [];
opts.regularizer = 'dwt';

opts = vararg_pair(opts,varargin);

sz_full = size(Y_GRAPPA); sz_full(end+1:4) = 1;

K = false(sz_full(1:3));
if isempty(opts.NFirsty), opts.NFirsty = mod(floor(sz_full(2)/2),Kernel.Ry)+1; end
if isempty(opts.NFirstz), opts.NFirstz = mod(floor(sz_full(3)/2),Kernel.Rz)+1; end
K(:,opts.NFirsty:Kernel.Ry:end,opts.NFirstz:Kernel.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

% perform GRAPPA + BayesDESIGN_SB with just acquired data
if isempty(opts.f_Y)
    switch lower(opts.regularizer)
        case 'dwt'
            f_Y = BayesDESIGN_SB(Y_GRAPPA,Lambda_G,cov_inds,sz_ACS,Kernel.Ry,Kernel.Rz,gamma,opts.BayesDESIGN_opts{:},'preservedata',opts.preservedata,'NFirsty',opts.NFirsty,'NFirstz',opts.NFirstz,'ACSFirstx',opts.ACSFirstx,'ACSFirsty',opts.ACSFirsty,'ACSFirstz',opts.ACSFirstz);
        case 'swt'
            f_Y = BayesDESIGN_SB_SWT(Y_GRAPPA,Lambda_G,cov_inds,sz_ACS,Kernel.Ry,Kernel.Rz,gamma,opts.BayesDESIGN_opts{:},'preservedata',opts.preservedata,'NFirsty',opts.NFirsty,'NFirstz',opts.NFirstz,'ACSFirstx',opts.ACSFirstx,'ACSFirsty',opts.ACSFirsty,'ACSFirstz',opts.ACSFirstz);
        case 'tv'
            f_Y = BayesDESIGN_SB_TV(Y_GRAPPA,Lambda_G,cov_inds,sz_ACS,Kernel.Ry,Kernel.Rz,gamma,opts.BayesDESIGN_opts{:},'preservedata',opts.preservedata,'NFirsty',opts.NFirsty,'NFirstz',opts.NFirstz,'ACSFirstx',opts.ACSFirstx,'ACSFirsty',opts.ACSFirsty,'ACSFirstz',opts.ACSFirstz);
        otherwise
            fail('unrecognized regularizer!');
    end
else
    f_Y = opts.f_Y;
end

% mask GRAPPA recon
Y_GRAPPA_mask = Y_GRAPPA;
f_Y_mask = f_Y;
if opts.preservedata
    Y_GRAPPA_mask(K(:,:,:,ones(1,sz_full(4)))) = 0;
    f_Y_mask(K(:,:,:,ones(1,sz_full(4)))) = 0;
end

% WSURE = ||B*f_gamma(y)||^2 - 2*real(y'*A'*B*f_gamma(y))
WSURE = abs(f_Y_mask(:)'*f_Y_mask(:)) - 2*real(Y_GRAPPA_mask(:)'*f_Y_mask(:));

% do the same with data + complex Bernoulli noise (GRAPPA is linear in data)
data_noise = b(:,opts.NFirsty:Kernel.Ry:end,opts.NFirstz:Kernel.Rz:end,:);
if isempty(sz_ACS)
    ACS_noise = [];
else
    ACS_noise = b(opts.ACSFirstx+(0:NACSx-1),opts.ACSFirsty+(0:NACSy-1),opts.ACSFirstz+(0:NACSz-1),:);
end
[Y_GRAPPA_noise,cov_inds] = GRAPPA_recon(sz_full,data_noise,Kernel,ACS_noise,opts.NFirsty,opts.NFirstz,opts.ACSFirstx,opts.ACSFirsty,opts.ACSFirstz);
switch lower(opts.regularizer)
    case 'dwt'
        f_Y_plus_b = BayesDESIGN_SB(Y_GRAPPA+epsilon.*Y_GRAPPA_noise,Lambda_G,cov_inds,sz_ACS,Kernel.Ry,Kernel.Rz,gamma,opts.BayesDESIGN_opts{:},'preservedata',opts.preservedata,'NFirsty',opts.NFirsty,'NFirstz',opts.NFirstz,'ACSFirstx',opts.ACSFirstx,'ACSFirsty',opts.ACSFirsty,'ACSFirstz',opts.ACSFirstz);
    case 'swt'
        f_Y_plus_b = BayesDESIGN_SB_SWT(Y_GRAPPA+epsilon.*Y_GRAPPA_noise,Lambda_G,cov_inds,sz_ACS,Kernel.Ry,Kernel.Rz,gamma,opts.BayesDESIGN_opts{:},'preservedata',opts.preservedata,'NFirsty',opts.NFirsty,'NFirstz',opts.NFirstz,'ACSFirstx',opts.ACSFirstx,'ACSFirsty',opts.ACSFirsty,'ACSFirstz',opts.ACSFirstz);
    case 'tv'
        f_Y_plus_b = BayesDESIGN_SB_TV(Y_GRAPPA+epsilon.*Y_GRAPPA_noise,Lambda_G,cov_inds,sz_ACS,Kernel.Ry,Kernel.Rz,gamma,opts.BayesDESIGN_opts{:},'preservedata',opts.preservedata,'NFirsty',opts.NFirsty,'NFirstz',opts.NFirstz,'ACSFirstx',opts.ACSFirstx,'ACSFirsty',opts.ACSFirsty,'ACSFirstz',opts.ACSFirstz);
    otherwise
        fail('unrecognized regularizer!');
end

varrho = (f_Y_plus_b - f_Y)./epsilon;

% mask the output (if preserving data)
if opts.preservedata
    varrho(K(:,:,:,ones(1,sz_full(4)))) = 0;
end

% mask the output (if preserving data)
Y_GRAPPA_b_mask = Y_GRAPPA_b;
if opts.preservedata
    Y_GRAPPA_b_mask(K(:,:,:,ones(1,sz_full(4)))) = 0;
end

WSURE = WSURE + 2*real(Y_GRAPPA_b_mask(:)'*varrho(:));

end
