% make_BayesDESIGN_plots.m -- make plots in Figs. 1, 2, 4, and
% 5 in MRM paper and Figs. S1, S3 in the supplement

% Copyright 2014-07-22 Daniel Weller, University of Michigan. All rights
% reserved.

% update path to image files
imfiles = {'/y/wellerds/MRI_data/MCSURE_v1.0/real1.mat',...
    '/y/wellerds/MRI_data/MCSURE_v1.0/real2.mat'};
imfile = 1; % for real image #1
imfile = 2; % for real image #2

if ~exist('figures','var')
%     figures = 1; % just make Fig. 1
%     figures = 2; % just make Fig. 2
%     figures = 4; % just make Fig. 4
%     figures = 5; % just make Fig. 5
%     figures = 9; % just make Fig. S1
%     figures = 11; % just make Fig. S3
    
    figures = 1:12; % make all figures
end

switch imfile
    case 1
        amp_factor_figs29 = 6; 
        amp_factor_figs411 = [48,96,192,384,768];
        amp_factor_fig5 = 96;

        window_image = [0,0.5];
        window_diff = [0,0.1];
        irect = [69,97,72,64];
        
        lambda = 1e4;
    case 2
        amp_factor_figs29 = 4; 
        amp_factor_figs411 = [8,16,32,64,128];
        amp_factor_fig5 = 8;

        window_image = [0,4];
        window_diff = [0,0.8];
        irect = [129,29,84,63];
        
        lambda = 1e2;
    otherwise
        error('Bad image #');
end

rng('default');

%% start pool
if any(figures == 2) || any(figures == 9)
    matlabpool('open','local');
end

%% load 8-channel coil sensitivities and brain slice from file 
load(imfiles{imfile},'ims','noisecov','mask');

% re-cast as doubles (seems to be useful for proper BayesDESIGN/CG behavior
ims = double(ims);
noisecov = double(noisecov);

% reshape
ims = permute(ims,[5,1,2,3,4]);
mask = permute(mask,[3,1,2]);

% enable/disable sensitivity normalization
smap_norm = [];

% get reference image
[Nx,Ny,Nz,NCha] = size(ims);
k_ref = fft3c(ims);
im_ref = sos_combine(ims,[],smap_norm);
mag = 100*max(size(im_ref)./[1,irect(3:4)]);

noisecov_init = noisecov;
clear ims;

%% generate plot of image, inset -- Fig. 1
if any(figures == 1)
    figure; imshow(squeeze(im_ref),window_image,'Border','tight'); hold on;
    hrect = rectangle('Position',irect([2,1,4,3]),'LineWidth',2,'EdgeColor','w');

    figure; imshow(squeeze(im_ref(:,irect(1) + (0:irect(3)-1),irect(2) + (0:irect(4)-1))),window_image,'Border','tight','InitialMagnification',mag);
end

%% get amplitude factors
amp_factors = [];
if any(figures == 2) || any(figures == 9)
    amp_factors = union(amp_factors,amp_factor_figs29);
end
if any(figures == 4) || any(figures == 11)
    amp_factors = union(amp_factors,amp_factor_figs411);
end
if any(figures == 5)
    amp_factors = union(amp_factors,amp_factor_fig5);
end

%% start simulating 
WSUREs = cell(size(amp_factors));
Y_GRAPPA_ref_normsqds = cell(size(amp_factors));
gammas = cell(size(amp_factors));
OracleMSEs = cell(size(amp_factors));
WMSEs = cell(size(amp_factors));
Ks = cell(size(amp_factors));
noisecovs = cell(size(amp_factors));
gamma_OMSEs = cell(size(amp_factors));
gamma_WSUREs = cell(size(amp_factors));
gamma_lcurves = cell(size(amp_factors));
Y_GRAPPAs = cell(size(amp_factors));
im_GRAPPAs = cell(size(amp_factors));
im_GRAPPA_diffs = cell(size(amp_factors));
Y_DESIGN_OMSEs = cell(size(amp_factors));
im_DESIGN_OMSEs = cell(size(amp_factors));
im_DESIGN_OMSE_diffs = cell(size(amp_factors));
Y_DESIGN_WSUREs = cell(size(amp_factors));
im_DESIGN_WSUREs = cell(size(amp_factors));
im_DESIGN_WSURE_diffs = cell(size(amp_factors));
Y_DESIGN_lcurves = cell(size(amp_factors));
im_DESIGN_lcurves = cell(size(amp_factors));
im_DESIGN_lcurve_diffs = cell(size(amp_factors));

rseed = rng();

% repeat over range of amplification factors
for iamp = 1:length(amp_factors)
    rng(rseed);
    amp_factor = amp_factors(iamp);
    
    %% amplify noise (optional)
    if amp_factor > 1
        noise = sqrtm((amp_factor-1).*noisecov_init/2)*complex(randn([NCha,Nx*Ny*Nz]),randn([NCha,Nx*Ny*Nz]));
        k_full = k_ref + reshape(noise.',[Nx,Ny,Nz,NCha]);
        noisecov = amp_factor.*noisecov_init;
    else
        noisecov = noisecov_init;
        k_full = k_ref;
    end
    im_full = sos_combine(ifft3c(k_full),[],smap_norm);
    noisecovs{iamp} = noisecov;

    %% undersample k-space
    Ry = 2;
    Rz = 2;
    NFirsty = mod(floor(Ny/2),Ry)+1;
    NFirstz = mod(floor(Nz/2),Rz)+1;
    K = false(Nx,Ny,Nz);
    k_redu = k_full(:,NFirsty:Ry:end,NFirstz:Rz:end,:);
    K(:,NFirsty:Ry:end,NFirstz:Rz:end) = true;

    % ACS lines for GRAPPA calibration
    NACSx = Nx; NACSy = 24; NACSz = 24;
    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_ACS = k_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;
    sz_ACS = [NACSx,NACSy,NACSz,NCha];
    Ks{iamp} = K;

    %% perform GRAPPA calibration
    Bx = 1; By = 4; Bz = 4; lambda_tk = 1e-2;
    [Kernel,Lambda_G] = GRAPPA_construct_kernel(k_ACS,Ry,Rz,Bx,By,Bz,noisecov,lambda_tk);

    %% do GRAPPA recon of noisy and reference k-space
    Y_GRAPPA_ref = GRAPPA_recon([Nx,Ny,Nz,NCha],k_ref(:,NFirsty:Ry:end,NFirstz:Rz:end,:),Kernel,k_ref(ACSFirstx:ACSFirstx+NACSx-1,ACSFirsty:ACSFirsty+NACSy-1,ACSFirstz:ACSFirstz+NACSz-1,:),NFirsty,NFirstz,ACSFirstx,ACSFirsty,ACSFirstz);
    Y_GRAPPA_ref_normsqd = Y_GRAPPA_ref;
    Y_GRAPPA_ref_normsqd(K(:,:,:,ones(1,NCha))) = 0;
    Y_GRAPPA_ref_normsqd = reshape(Y_GRAPPA_ref_normsqd,[Nx*Ny*Nz,NCha]);
    Y_GRAPPA_ref_normsqd = abs(Y_GRAPPA_ref_normsqd(:)'*Y_GRAPPA_ref_normsqd(:));
    Y_GRAPPA_ref_normsqds{iamp} = Y_GRAPPA_ref_normsqd;

    epsilon_scale_factor = norm(col(k_redu))/sqrt(numel(k_redu));
    [Y_GRAPPA,cov_inds] = GRAPPA_recon([Nx,Ny,Nz,NCha],k_redu,Kernel,k_ACS,NFirsty,NFirstz,ACSFirstx,ACSFirsty,ACSFirstz);
    im_GRAPPA = sos_combine(ifft3c(Y_GRAPPA),[],smap_norm);
    im_GRAPPA_diff = im_GRAPPA-im_ref;
    Y_GRAPPAs{iamp} = Y_GRAPPA;
    im_GRAPPAs{iamp} = im_GRAPPA;
    im_GRAPPA_diffs{iamp} = im_GRAPPA_diff;

    %% BayesDESIGN options
    BayesDESIGN_opts = {'lambda',lambda,'maxiters',20,'forcemaxiters',true,'tv_isotropic',true};

    %% generate noise for WSURE estimation
    if (any(figures == 2) || any(figures == 9)) && any(amp_factor_figs29 == amp_factor)
        Nbs = 20;
        bs = complex(zeros([Nx,Ny,Nz,NCha,Nbs]));
        Y_GRAPPA_bs = complex(zeros([Nx,Ny,Nz,NCha,Nbs]));
        Kcha = K(:,:,:,ones(1,NCha));
        parfor ib = 1:Nbs
            breal = rand([Nx,Ny,Nz,NCha]);
            breal(breal >= 0.5) = 1;
            breal(breal < 0.5) = -1;
            bimag = rand([Nx,Ny,Nz,NCha]);
            bimag(bimag >= 0.5) = 1;
            bimag(bimag < 0.5) = -1;
            b = complex(breal,bimag)./sqrt(2);
            b(~Kcha) = 0;
            bs(:,:,:,:,ib) = b;

            % compute SPIRiT(Lambda*b)
            Y_GRAPPA_b = reshape((noisecov*reshape(b,[Nx*Ny*Nz,NCha]).').',[Nx,Ny,Nz,NCha]);
            Y_GRAPPA_b = GRAPPA_recon([Nx,Ny,Nz,NCha],Y_GRAPPA_b(:,NFirsty:Ry:end,NFirstz:Rz:end,:),Kernel,Y_GRAPPA_b(ACSFirstx:ACSFirstx+NACSx-1,ACSFirsty:ACSFirsty+NACSy-1,ACSFirstz:ACSFirstz+NACSz-1,:),NFirsty,NFirstz,ACSFirstx,ACSFirsty,ACSFirstz);
            Y_GRAPPA_bs(:,:,:,:,ib) = Y_GRAPPA_b;
        end
        b = bs(:,:,:,:,1);
        Y_GRAPPA_b = Y_GRAPPA_bs(:,:,:,:,1);
    else
        Kcha = K(:,:,:,ones(1,NCha));
        breal = rand([Nx,Ny,Nz,NCha]);
        breal(breal >= 0.5) = 1;
        breal(breal < 0.5) = -1;
        bimag = rand([Nx,Ny,Nz,NCha]);
        bimag(bimag >= 0.5) = 1;
        bimag(bimag < 0.5) = -1;
        b = complex(breal,bimag)./sqrt(2);
        b(~Kcha) = 0;

        % compute SPIRiT(Lambda*b)
        Y_GRAPPA_b = reshape((noisecov*reshape(b,[Nx*Ny*Nz,NCha]).').',[Nx,Ny,Nz,NCha]);
        Y_GRAPPA_b = GRAPPA_recon([Nx,Ny,Nz,NCha],Y_GRAPPA_b(:,NFirsty:Ry:end,NFirstz:Rz:end,:),Kernel,Y_GRAPPA_b(ACSFirstx:ACSFirstx+NACSx-1,ACSFirsty:ACSFirsty+NACSy-1,ACSFirstz:ACSFirstz+NACSz-1,:),NFirsty,NFirstz,ACSFirstx,ACSFirsty,ACSFirstz);
    end

    %% compare WSURE, WMSE, OMSE values for Fig. 2
    if (any(figures == 2) || any(figures == 9)) && any(amp_factor_figs29 == amp_factor)
        epsilons = logspace(-5,-2,4); % try a few to get an idea of what is small enough
        switch imfile 
            case 1
                gammas1 = union(logspace(1,4,13),logspace(2,3,9),'sorted');
            case 2
                gammas1 = union(logspace(0,3,13),logspace(1,2,9),'sorted'); % use smaller gammas here
        end
            
        WSUREs1 = zeros(length(gammas1),length(epsilons),Nbs);
        WMSEs1 = zeros(length(gammas1),1);
        OracleMSEs1 = zeros(length(gammas1),1);
        fprintf(1,'Comparing WSURE, WMSE, OMSE values for various gammas, epsilons...\n'); drawnow;
        for ig = 1:length(gammas1)
            gamma = gammas1(ig);
            for ie = 1:length(epsilons)
                epsilon = epsilons(ie);
                if ie == 1
                    [WSUREs1(ig,ie,1),f_Y] = BayesDESIGN_SB_WSURE3(Y_GRAPPA,cov_inds,bs(:,:,:,:,1),Y_GRAPPA_bs(:,:,:,:,1),Kernel,Lambda_G,sz_ACS,gamma,epsilon*epsilon_scale_factor,'NFirsty',NFirsty,'NFirstz',NFirstz,'ACSFirstx',ACSFirstx,'ACSFirsty',ACSFirsty,'ACSFirstz',ACSFirstz,'preservedata',true,'BayesDESIGN_opts',BayesDESIGN_opts,'regularizer','tv');
                    parfor ib = 2:Nbs
                        WSUREs1(ig,ie,ib) = BayesDESIGN_SB_WSURE3(Y_GRAPPA,cov_inds,bs(:,:,:,:,ib),Y_GRAPPA_bs(:,:,:,:,ib),Kernel,Lambda_G,sz_ACS,gamma,epsilon*epsilon_scale_factor,'NFirsty',NFirsty,'NFirstz',NFirstz,'ACSFirstx',ACSFirstx,'ACSFirsty',ACSFirsty,'ACSFirstz',ACSFirstz,'preservedata',true,'BayesDESIGN_opts',BayesDESIGN_opts,'f_Y',f_Y,'regularizer','tv');
                    end
                else
                    parfor ib = 1:Nbs
                        WSUREs1(ig,ie,ib) = BayesDESIGN_SB_WSURE3(Y_GRAPPA,cov_inds,bs(:,:,:,:,ib),Y_GRAPPA_bs(:,:,:,:,ib),Kernel,Lambda_G,sz_ACS,gamma,epsilon*epsilon_scale_factor,'NFirsty',NFirsty,'NFirstz',NFirstz,'ACSFirstx',ACSFirstx,'ACSFirsty',ACSFirsty,'ACSFirstz',ACSFirstz,'preservedata',true,'BayesDESIGN_opts',BayesDESIGN_opts,'f_Y',f_Y,'regularizer','tv');
                    end
                end
                fprintf(1,' %9g',mean(WSUREs1(ig,ie,:))+Y_GRAPPA_ref_normsqd); drawnow;
            end
            OMSE = f_Y - k_ref;
            OMSE(K(:,:,:,ones(1,NCha))) = 0;
            OMSE = abs(OMSE(:)'*OMSE(:));
            OracleMSEs1(ig) = OMSE;
            WMSE = f_Y - Y_GRAPPA_ref;
            WMSE(K(:,:,:,ones(1,NCha))) = 0;
            WMSE = abs(WMSE(:)'*WMSE(:));
            WMSEs1(ig) = WMSE;
            fprintf(1,' %9g %9g\n',WMSEs1(ig),OracleMSEs1(ig)); drawnow;
        end
        WSUREs{iamp} = WSUREs1;
        gammas{iamp} = gammas1;
        OracleMSEs{iamp} = OracleMSEs1;
        WMSEs{iamp} = WMSEs1;
    end

    %% done with multiple-noise part
    if (any(figures == 2) || any(figures == 9)) && any(amp_factor_figs29 == amp_factor)
        clear bs Y_GRAPPA_bs;
    end

    %% do brute-force search for optimal gamma that minimizes true MSE
    if (any(figures == 4) || any(figures == 5) || any(figures == 11)) && (any(amp_factor_figs411 == amp_factor) || any(amp_factor_fig5 == amp_factor))
        tDESIGN_OMSE = tic();
        fprintf(1,'Using fminsearch() to find Oracle WMSE-optimal gamma...\n'); drawnow;
        switch imfile
            case 1
                loggamma_opt = log10(1e3);
            case 2
                loggamma_opt = log10(10);
        end
        [loggamma_opt,OMSE_opt,~,fminsearch_output] = fminsearch(@(loggamma) sum(sum(abs(reshape((~K(:,:,:,ones(1,NCha))).*(BayesDESIGN_SB_TV(Y_GRAPPA,Lambda_G,cov_inds,sz_ACS,Ry,Rz,10.^loggamma,'NFirsty',NFirsty,'NFirstz',NFirstz,'ACSFirstx',ACSFirstx,'ACSFirsty',ACSFirsty,'ACSFirstz',ACSFirstz,'preservedata',true,BayesDESIGN_opts{:}) - k_ref),[Nx*Ny*Nz,NCha])).^2)),loggamma_opt,optimset('Display','iter','FunValCheck','on'));
        gamma_OMSE = 10.^loggamma_opt;
        tDESIGN_OMSE = toc(tDESIGN_OMSE);
        DESIGNevals_OMSE = fminsearch_output.funcCount;
        fprintf(1,'Oracle WMSE search: min(OMSE) = %g at gamma = %g.\n',OMSE_opt,gamma_OMSE); drawnow;
        Y_DESIGN_OMSE = BayesDESIGN_SB_TV(Y_GRAPPA,Lambda_G,cov_inds,sz_ACS,Ry,Rz,gamma_OMSE,'NFirsty',NFirsty,'NFirstz',NFirstz,'ACSFirstx',ACSFirstx,'ACSFirsty',ACSFirsty,'ACSFirstz',ACSFirstz,'preservedata',true,BayesDESIGN_opts{:});
        im_DESIGN_OMSE = sos_combine(ifft3c(Y_DESIGN_OMSE),[],smap_norm);
        im_DESIGN_OMSE_diff = im_DESIGN_OMSE-im_ref;
        gamma_OMSEs{iamp} = gamma_OMSE;
        Y_DESIGN_OMSEs{iamp} = Y_DESIGN_OMSE;
        im_DESIGN_OMSEs{iamp} = im_DESIGN_OMSE;
        im_DESIGN_OMSE_diffs{iamp} = im_DESIGN_OMSE_diff;
    end

    %% do brute-force search for optimal gamma that minimizes WSURE
    if (any(figures == 4) || any(figures == 5) || any(figures == 11)) && (any(amp_factor_figs411 == amp_factor) || any(amp_factor_fig5 == amp_factor))
        tDESIGN_WSURE = tic();
        epsilon = 1e-4; % base on earlier evaluations for what works best
        Ngammas = 9;
        switch imfile
            case 1
                loggamma_range = [-3,7];
            case 2
                loggamma_range = [-4,6];
        end
        Nlevels = 2;
        WSUREs_search = zeros(Nlevels,Ngammas);
        loggammas = zeros(Nlevels,Ngammas+2);
        fprintf(1,'Using coarse-to-fine parameter sweep to find WSURE-optimal gamma...\n'); drawnow;
        for ilvl = 1:Nlevels
            loggammas(ilvl,:) = linspace(loggamma_range(1),loggamma_range(2),Ngammas+2);
            for ig = 1:Ngammas
                WSUREs_search(ilvl,ig) = BayesDESIGN_SB_WSURE3(Y_GRAPPA,cov_inds,b,Y_GRAPPA_b,Kernel,Lambda_G,sz_ACS,10.^loggammas(ilvl,ig+1),epsilon*epsilon_scale_factor,'NFirsty',NFirsty,'NFirstz',NFirstz,'ACSFirstx',ACSFirstx,'ACSFirsty',ACSFirsty,'ACSFirstz',ACSFirstz,'preservedata',true,'BayesDESIGN_opts',BayesDESIGN_opts,'regularizer','tv');
                fprintf(1,'WSURE search level %d: WSURE(%g) = %g + ||GMx||^2.\n',ilvl,10.^loggammas(ilvl,ig+1),WSUREs_search(ilvl,ig)); drawnow;
            end
            [WSURE_opt,ig] = min(WSUREs_search(ilvl,:));
            gamma_WSURE = 10.^loggammas(ilvl,ig+1);
            fprintf(1,'WSURE search level %d: min(WSURE) = %g + ||GMx||^2 at gamma = %g.\n',ilvl,WSURE_opt,gamma_WSURE); drawnow;
            loggamma_range = loggammas(ilvl,[ig,ig+2]);
        end
        tDESIGN_WSURE = toc(tDESIGN_WSURE);
        DESIGNevals_WSURE = 2*Ngammas*Nlevels;
        Y_DESIGN_WSURE = BayesDESIGN_SB_TV(Y_GRAPPA,Lambda_G,cov_inds,sz_ACS,Ry,Rz,gamma_WSURE,'NFirsty',NFirsty,'NFirstz',NFirstz,'ACSFirstx',ACSFirstx,'ACSFirsty',ACSFirsty,'ACSFirstz',ACSFirstz,'preservedata',true,BayesDESIGN_opts{:});
        im_DESIGN_WSURE = sos_combine(ifft3c(Y_DESIGN_WSURE),[],smap_norm);
        im_DESIGN_WSURE_diff = im_DESIGN_WSURE-im_ref;
        gamma_WSUREs{iamp} = gamma_WSURE;
        Y_DESIGN_WSUREs{iamp} = Y_DESIGN_WSURE;
        im_DESIGN_WSUREs{iamp} = im_DESIGN_WSURE;
        im_DESIGN_WSURE_diffs{iamp} = im_DESIGN_WSURE_diff;
    end

    %% do "L-curve" type of analysis
    if (any(figures == 4) || any(figures == 5) || any(figures == 11)) && (any(amp_factor_figs411 == amp_factor) || any(amp_factor_fig5 == amp_factor))
        tDESIGN_Lcurve = tic();
        Ngammas = 9;
        switch imfile
            case 1
                loggamma_range = [-3,7];
            case 2
                loggamma_range = [-4,6];
        end
        Nlevels = 2;
        L1objs = zeros(Nlevels,Ngammas);
        GRAPPAobjs = zeros(Nlevels,Ngammas);
        loggammas = zeros(Nlevels,Ngammas+2);
        fprintf(1,'Using coarse-to-fine parameter sweep to do l-curve analysis...\n'); drawnow;
        [~,~,GRAPPAobj_min,L1obj_min] = BayesDESIGN_SB_TV(Y_GRAPPA,Lambda_G,cov_inds,sz_ACS,Ry,Rz,10.^loggamma_range(1),'NFirsty',NFirsty,'NFirstz',NFirstz,'ACSFirstx',ACSFirstx,'ACSFirsty',ACSFirsty,'ACSFirstz',ACSFirstz,'preservedata',true,BayesDESIGN_opts{:});
        [~,~,GRAPPAobj_max,L1obj_max] = BayesDESIGN_SB_TV(Y_GRAPPA,Lambda_G,cov_inds,sz_ACS,Ry,Rz,10.^loggamma_range(2),'NFirsty',NFirsty,'NFirstz',NFirstz,'ACSFirstx',ACSFirstx,'ACSFirsty',ACSFirsty,'ACSFirstz',ACSFirstz,'preservedata',true,BayesDESIGN_opts{:});
        for ilvl = 1:Nlevels
            loggammas(ilvl,:) = linspace(loggamma_range(1),loggamma_range(2),Ngammas+2);
            for ig = 1:Ngammas
                [~,~,GRAPPAobjs(ilvl,ig),L1objs(ilvl,ig)] = BayesDESIGN_SB_TV(Y_GRAPPA,Lambda_G,cov_inds,sz_ACS,Ry,Rz,10.^loggammas(ilvl,ig+1),'NFirsty',NFirsty,'NFirstz',NFirstz,'ACSFirstx',ACSFirstx,'ACSFirsty',ACSFirsty,'ACSFirstz',ACSFirstz,'preservedata',true,BayesDESIGN_opts{:});
                fprintf(1,'L-curve search level %d: GRAPPAobj(%g) = %g, L1obj(%g) = %g.\n',ilvl,10.^loggammas(ilvl,ig+1),GRAPPAobjs(ilvl,ig),10.^loggammas(ilvl,ig+1),L1objs(ilvl,ig)); drawnow;
            end
            L1deriv_post = (log10([L1objs(ilvl,2:end),L1obj_max(ilvl)]) - log10(L1objs(ilvl,:)));
            L1deriv_pre = (log10(L1objs(ilvl,:)) - log10([L1obj_min(ilvl),L1objs(ilvl,1:end-1)]));
            GRAPPAderiv_post = log10([GRAPPAobjs(ilvl,2:end),GRAPPAobj_max(ilvl)]) - log10(GRAPPAobjs(ilvl,:));
            GRAPPAderiv_pre = log10(GRAPPAobjs(ilvl,:)) - log10([GRAPPAobj_min(ilvl),GRAPPAobjs(ilvl,1:end-1)]);

            deriv = (L1deriv_post + L1deriv_pre)./(GRAPPAderiv_post + GRAPPAderiv_pre);
            deriv2 = 4.*(L1deriv_post-L1deriv_pre)./(GRAPPAderiv_post + GRAPPAderiv_pre).^2;

            firstvalidderiv = find(deriv <= 0 & ~isnan(deriv),1,'first');
            if isempty(firstvalidderiv), firstvalidderiv = length(deriv)+1; end
            deriv2 = deriv2(firstvalidderiv:end);
            firstinvalidderiv = find(deriv(firstvalidderiv:end) > 0 | isnan(deriv(firstvalidderiv:end)),1,'first');
            if ~isempty(firstinvalidderiv)
                deriv2 = deriv2(1:firstinvalidderiv-1);
                if firstinvalidderiv > 1
    %                 deriv2(firstinvalidderiv-1) = -L1deriv_pre(firstinvalidderiv-1)./GRAPPAderiv_pre(firstinvalidderiv-1).^2;
                    deriv2 = deriv2(1:firstinvalidderiv-2);
                end
            end
            if isempty(deriv2)
                fprintf(2,'L-curve search level %d: no valid part of pareto curve found.\n',ilvl); drawnow;
                deriv2_max = NaN; ig = floor(Ngammas/2)+1;
%                 break;
            else
                [deriv2_max,ig] = max(deriv2);
                ig = ig + (firstvalidderiv-1);
                if deriv2_max <= 0
                    fprintf(1,'L-curve search level %d: pareto curve seems to be concave. Continuing...\n',ilvl); drawnow;
                end
            end

            gamma_lcurve = 10.^loggammas(ilvl,ig+1);
            fprintf(1,'L-curve search level %d: max(deriv2) = %g at gamma = %g.\n',ilvl,deriv2_max,gamma_lcurve); drawnow;
            if ig == 1, L1obj_min(ilvl+1) = L1obj_min(ilvl); else L1obj_min(ilvl+1) = L1objs(ilvl,ig-1); end
            if ig == Ngammas, L1obj_max(ilvl+1) = L1obj_max(ilvl); else L1obj_max(ilvl+1) = L1objs(ilvl,ig+1); end
            if ig == 1, GRAPPAobj_min(ilvl+1) = GRAPPAobj_min(ilvl); else GRAPPAobj_min(ilvl+1) = GRAPPAobjs(ilvl,ig-1); end
            if ig == Ngammas, GRAPPAobj_max(ilvl+1) = GRAPPAobj_max(ilvl); else GRAPPAobj_max(ilvl+1) = GRAPPAobjs(ilvl,ig+1); end
            loggamma_range = loggammas(ilvl,[ig,ig+2]);
        end
        tDESIGN_Lcurve = toc(tDESIGN_Lcurve);
        DESIGNevals_Lcurve = Ngammas*Nlevels+2;
        Y_DESIGN_lcurve = BayesDESIGN_SB_TV(Y_GRAPPA,Lambda_G,cov_inds,sz_ACS,Ry,Rz,gamma_lcurve,'NFirsty',NFirsty,'NFirstz',NFirstz,'ACSFirstx',ACSFirstx,'ACSFirsty',ACSFirsty,'ACSFirstz',ACSFirstz,'preservedata',true,BayesDESIGN_opts{:});
        im_DESIGN_lcurve = sos_combine(ifft3c(Y_DESIGN_lcurve),[],smap_norm);
        im_DESIGN_lcurve_diff = im_DESIGN_lcurve-im_ref;
        gamma_lcurves{iamp} = gamma_lcurve;
        Y_DESIGN_lcurves{iamp} = Y_DESIGN_lcurve;
        im_DESIGN_lcurves{iamp} = im_DESIGN_lcurve;
        im_DESIGN_lcurve_diffs{iamp} = im_DESIGN_lcurve_diff;
    end

%% end simulations
end

%% cleanup

if any(figures == 2) || any(figures == 9)
    matlabpool('close');
end

%% do plots

%% WSURE vs. true and approx. WMSE plot -- Fig. 2

if any(figures == 2)
    iamp = find(amp_factors == amp_factor_figs29(1),1);

    % compute WSURE stuff
    eps_index = 2; % epsilon = 1e-4
    nstds = 5;
    if isempty(eps_index), error('Desired epsilon not found.'); end
    WSURE_val = WSUREs{iamp}(:,eps_index,1) + Y_GRAPPA_ref_normsqds{iamp};
    WSURE_mean = mean(WSUREs{iamp}(:,eps_index,:),3) + Y_GRAPPA_ref_normsqds{iamp};
    WSURE_std = std(WSUREs{iamp}(:,eps_index,:),0,3);
    WSURE_lower = WSURE_mean - nstds.*WSURE_std;
    WSURE_upper = WSURE_mean + nstds.*WSURE_std;

    % make plot
    figure; 
%     subplot(2,1,1); 
    plot(log10(gammas{iamp}),10*log10(OracleMSEs{iamp}),'-o',log10(gammas{iamp}),10*log10(WMSEs{iamp}),'-s'); hold all;
    errorbar(log10(gammas{iamp}),10*log10(WSURE_val),10*log10(WSURE_val./WSURE_lower),10*log10(WSURE_upper./WSURE_val));
    hcs = get(gca,'Children');
    legend(hcs([3,1,2]),'True WMSE (Eq. [3])','WSURE (\epsilon = 10^{-4})','Approx. WMSE (Eq. [4])','Location','NorthWest');
    xlim([floor(log10(min(gammas{iamp}))),ceil(log10(max(gammas{iamp})))]);
    xticks = floor(log10(min(gammas{iamp}))):ceil(log10(max(gammas{iamp})));
    xtickstrs = arrayfun(@(x) sprintf('%g',10.^x),xticks,'UniformOutput',false);
    set(gca,'Box','off','XTick',xticks,'XTickLabel',xtickstrs);
    xlabel('Regularization parameter (\gamma)');
    ylabel('Squared error (dB)');
end

%% WSURE vs. true and approx. WMSE plot for all epsilons -- Fig. S1

if any(figures == 9)
    iamp = find(amp_factors == amp_factor_figs29(1),1);

    % compute WSURE stuff
    WSURE_vals = WSUREs{iamp}(:,:,1) + Y_GRAPPA_ref_normsqds{iamp};
    WSURE_means = mean(WSUREs{iamp},3) + Y_GRAPPA_ref_normsqds{iamp};
    WSURE_stds = std(WSUREs{iamp},0,3);
    WSURE_lowers = WSURE_means - WSURE_stds;
    WSURE_uppers = WSURE_means + WSURE_stds;

    % make plot
    figure; 
%     subplot(2,1,1); 
    plot(log10(gammas{iamp}),10*log10(OracleMSEs{iamp}),'-o',log10(gammas{iamp}),10*log10(WMSEs{iamp}),'-s',log10(gammas{iamp}),10*log10(WSURE_vals),'--');
    hcs = get(gca,'Children');
    legendstrs = {'True WMSE (Eq. [3])','WSURE (\epsilon = 10^{-5})','WSURE (\epsilon = 10^{-4})','WSURE (\epsilon = 0.001)','WSURE (\epsilon = 0.01)','Approx. WMSE (Eq. [4])'};
    legend(hcs([6,4:-1:1,5]),legendstrs,'Location','NorthWest');
    xlim([floor(log10(min(gammas{iamp}))),ceil(log10(max(gammas{iamp})))]);
    xticks = floor(log10(min(gammas{iamp}))):ceil(log10(max(gammas{iamp})));
    xtickstrs = arrayfun(@(x) sprintf('%g',10.^x),xticks,'UniformOutput',false);
    set(gca,'Box','off','XTick',xticks,'XTickLabel',xtickstrs);
    xlabel('Regularization parameter (\gamma)');
    ylabel('Squared error (dB)');
end

%% compute trend plot for WSURE-optimal and WMSE-optimal DESIGN recons vs. noise level -- Fig. 4 or Fig. S3

if any(figures == 4) || any(figures == 11)
    [~,iamps] = intersect(amp_factors,amp_factor_figs411);
    
    SNRs = zeros(size(iamps));
    gammas_DESIGN_OMSE = zeros(size(iamps));
    gammas_DESIGN_Lcurve = zeros(size(iamps));
    gammas_DESIGN_WSURE = zeros(size(iamps));
    WMSEs_GRAPPA = zeros(size(iamps));
    WMSEs_DESIGN_OMSE = zeros(size(iamps));
    WMSEs_DESIGN_Lcurve = zeros(size(iamps));
    WMSEs_DESIGN_WSURE = zeros(size(iamps));

    for ii = 1:length(iamps)
        iamp = iamps(ii);
        SNRs(ii) = norm(k_ref(:))^2/(numel(Ks{iamp})*sum(abs(svd(noisecovs{iamp}))));
        gammas_DESIGN_OMSE(ii) = gamma_OMSEs{iamp};
        gammas_DESIGN_Lcurve(ii) = gamma_lcurves{iamp};
        gammas_DESIGN_WSURE(ii) = gamma_WSUREs{iamp};
        KCha = Ks{iamp}(:,:,:,ones(1,size(noisecovs{iamp},1)));
        WMSEs_GRAPPA(ii) = norm((~KCha(:)).*(Y_GRAPPAs{iamp}(:)-k_ref(:)))^2;
        WMSEs_DESIGN_OMSE(ii) = norm((~KCha(:)).*(Y_DESIGN_OMSEs{iamp}(:)-k_ref(:)))^2;
        WMSEs_DESIGN_Lcurve(ii) = norm((~KCha(:)).*(Y_DESIGN_lcurves{iamp}(:)-k_ref(:)))^2;
        WMSEs_DESIGN_WSURE(ii) = norm((~KCha(:)).*(Y_DESIGN_WSUREs{iamp}(:)-k_ref(:)))^2;
    end

    sortedvecs = sortrows([SNRs(:),gammas_DESIGN_OMSE(:),gammas_DESIGN_Lcurve(:),gammas_DESIGN_WSURE(:),WMSEs_GRAPPA(:),WMSEs_DESIGN_OMSE(:),WMSEs_DESIGN_Lcurve(:),WMSEs_DESIGN_WSURE(:)],1);

    figure; subplot(2,1,2); plot(10*log10(sortedvecs(:,1)),sortedvecs(:,2),'-x',10*log10(sortedvecs(:,1)),sortedvecs(:,4),'-o',10*log10(sortedvecs(:,1)),sortedvecs(:,3),'-*');
    xlabel('Input SNR (dB)');
    ylabel('Regularization parameter (\gamma)');
    hcs = get(gca,'Children');
    set(hcs(3),'MarkerSize',14);
    xlim(10*log10(sortedvecs([1,end],1)));
    set(gca,'Box','off');

    subplot(2,1,1); plot(10*log10(sortedvecs(:,1)),10*log10(sortedvecs(:,6)),'-x',10*log10(sortedvecs(:,1)),10*log10(sortedvecs(:,8)),'-o',10*log10(sortedvecs(:,1)),10*log10(sortedvecs(:,7)),'-*',10*log10(sortedvecs(:,1)),10*log10(sortedvecs(:,5)),'-s');
    xlabel('Input SNR (dB)');
    ylabel('WMSE (dB)');
    hcs = get(gca,'Children');
    set(hcs(4),'MarkerSize',14);
    legend(hcs([1,2,3,4]),'GRAPPA','Lcurve-opt. DESIGN','WSURE-opt. DESIGN','WMSE-opt. DESIGN','Location','NorthEast');
    xlim(10*log10(sortedvecs([1,end],1)));
    set(gca,'Box','off');
end

%% show reconstructed images -- Fig. 5

if any(figures == 5)
    iamp = find(amp_factors == amp_factor_fig5(1),1);
    
    % GRAPPA alone
    figure; imshow(squeeze(im_GRAPPAs{iamp}),window_image,'Border','tight'); hold on;
    hin = axes('Position',[0.5,0,0.5,0.5]);
    imagesc(squeeze(im_GRAPPAs{iamp}(:,irect(1) + (0:irect(3)-1),irect(2) + (0:irect(4)-1))),window_image); colormap gray; axis image;
    set(hin,'XTick',[],'YTick',[],'Box','on','LineWidth',2);
    figure; imshow(squeeze(im_GRAPPA_diffs{iamp}),window_diff,'Border','tight'); hold on;
    hin = axes('Position',[0.5,0,0.5,0.5]);
    imagesc(squeeze(im_GRAPPA_diffs{iamp}(:,irect(1) + (0:irect(3)-1),irect(2) + (0:irect(4)-1))),window_diff); colormap gray; axis image;
    set(hin,'XTick',[],'YTick',[],'Box','on','LineWidth',2);

    % oracle
    figure; imshow(squeeze(im_DESIGN_OMSEs{iamp}),window_image,'Border','tight'); hold on;
    hin = axes('Position',[0.5,0,0.5,0.5]);
    imagesc(squeeze(im_DESIGN_OMSEs{iamp}(:,irect(1) + (0:irect(3)-1),irect(2) + (0:irect(4)-1))),window_image); colormap gray; axis image;
    set(hin,'XTick',[],'YTick',[],'Box','on','LineWidth',2);
    figure; imshow(squeeze(im_DESIGN_OMSE_diffs{iamp}),window_diff,'Border','tight'); hold on;
    hin = axes('Position',[0.5,0,0.5,0.5]);
    imagesc(squeeze(im_DESIGN_OMSE_diffs{iamp}(:,irect(1) + (0:irect(3)-1),irect(2) + (0:irect(4)-1))),window_diff); colormap gray; axis image;
    set(hin,'XTick',[],'YTick',[],'Box','on','LineWidth',2);

    % lcurve
    figure; imshow(squeeze(im_DESIGN_lcurves{iamp}),window_image,'Border','tight'); hold on;
    hin = axes('Position',[0.5,0,0.5,0.5]);
    imagesc(squeeze(im_DESIGN_lcurves{iamp}(:,irect(1) + (0:irect(3)-1),irect(2) + (0:irect(4)-1))),window_image); colormap gray; axis image;
    set(hin,'XTick',[],'YTick',[],'Box','on','LineWidth',2);
    figure; imshow(squeeze(im_DESIGN_lcurve_diffs{iamp}),window_diff,'Border','tight'); hold on;
    hin = axes('Position',[0.5,0,0.5,0.5]);
    imagesc(squeeze(im_DESIGN_lcurve_diffs{iamp}(:,irect(1) + (0:irect(3)-1),irect(2) + (0:irect(4)-1))),window_diff); colormap gray; axis image;
    set(hin,'XTick',[],'YTick',[],'Box','on','LineWidth',2);

    % WSURE
    figure; imshow(squeeze(im_DESIGN_WSUREs{iamp}),window_image,'Border','tight'); hold on;
    hcb = colorbar('Location','South');
    set(hcb,'XAxisLocation','bottom','XTick',window_image);
    hin = axes('Position',[0.5,0,0.5,0.5]);
    imagesc(squeeze(im_DESIGN_WSUREs{iamp}(:,irect(1) + (0:irect(3)-1),irect(2) + (0:irect(4)-1))),window_image); colormap gray; axis image;
    set(hin,'XTick',[],'YTick',[],'Box','on','LineWidth',2);
    figure; imshow(squeeze(im_DESIGN_WSURE_diffs{iamp}),window_diff,'Border','tight'); hold on;
    hcb = colorbar('Location','South');
    set(hcb,'XAxisLocation','bottom','XTick',window_diff);
    hin = axes('Position',[0.5,0,0.5,0.5]);
    imagesc(squeeze(im_DESIGN_WSURE_diffs{iamp}(:,irect(1) + (0:irect(3)-1),irect(2) + (0:irect(4)-1))),window_diff); colormap gray; axis image;
    set(hin,'XTick',[],'YTick',[],'Box','on','LineWidth',2);
end
