% make plots for L1SPIRiT ADMM T-MI paper using Brainweb data
%
% Code requires both SPIRiT v0.1 or v0.3, as well as IRT toolbox
%
% Copyright 2014-07-23, Daniel Weller, University of Michigan. All rights
% reserved. Email wellerds@umich.edu.

% choose a file to process (TODO: update paths)
imfiles = { ...
'../data/brainweb_8ch.mat'; ...
'../data/brainweb_16ch.mat'; ...
%'/y/wellerds/MRI_data/L1SPIRiT_ADMM_v1.0/brainweb_8ch.mat'; ...
%'/y/wellerds/MRI_data/L1SPIRiT_ADMM_v1.0/brainweb_16ch.mat' ...
};
imfile = 1; % 8-channel simulated data; used for calibrating mus
% imfile = 2; % 16-channel simulated data

mus_test = true; % set to true to run for a grid of mu's (Fig. 4 with image 1)
full_convergence = false; % set to true to get supplementary convergence plots; false gets plots in paper

% path initialization
rng('default');
if ~exist('SPIRiT', 'class')
%	addpath('../SPIRiT_v0.1/');
%	addpath('../ESPIRiT/');
	addpath('~/l/dat/code/matlab/ESPIRiT/');
end
if ~exist('zpad.m', 'file')
%	addpath('../SPIRiT_v0.1/utils/');
%	addpath('../ESPIRiT/utils/');
	addpath('~/l/dat/code/matlab/ESPIRiT/utils');
end
if ~exist('ticker.m', 'file')
	error('code requires IRT to be in path');
end

switch imfile
    case 1
        accel = 3;
        CalibSize = [30,30];            % size of the calibration region
        tStore = 50; % in seconds

        irect = [137,137,64,64];
        window_image = [0,1];
        window_diff_scale = 10;
        flipimfun = @(im) im; % no flip
    case 2
        accel = 4;
        CalibSize = [36,36];            % size of the calibration region
        tStore = 100; % in seconds

        irect = [137,137,64,64];
        window_image = [0,1];
        window_diff_scale = 5;
        flipimfun = @(im) im; % no flip
    otherwise
        error('unknown file selected.');
end

baseline_algorithm = 'SB'; % use 'ADMM' for faster baseline, 'SB' for paper

%% load BrainMap data with coil sensitivities
load(imfiles{imfile}, 'im_ref', 'smap', 'noisecov');

im_ref = double(im_ref);
smap = double(smap);
[Nx,Ny,NCha] = size(smap);

%% make multi-channel images
im_ref = im_ref(:,:,ones(1,NCha)).*smap;
smap_norm = sos(smap);
for ich = 1:NCha
    im(:,:,ich) = complex(BsplReduce(single(real(im_ref(:,:,ich))),3),BsplReduce(single(imag(im_ref(:,:,ich))),3)); %#ok<SAGROW>
end
im = double(im);
im_sqr = sos(im);
clear smap; % don't need sensitivities anymore

nrmsefun = @(x) norm(col(x - im))/norm(col(im));
mag = 100*max(size(im_sqr)./irect(3:4));

%% perform undersampling with random spiral trajectory

interleaves = 120;
N = min(Nx,Ny)/2;

% function for computing spiral trajectory
[~,omegas] = mri_kspace_spiral('N',N,'nl',interleaves,'Nt',0);

% undersample by selecting interleaves
omegas = cellfun(@(x) reshape(x,[],interleaves),num2cell(omegas,1),'UniformOutput',false);
rleaves = sort(randperm(interleaves,ceil(interleaves/accel))); % random leaves
omegas = cellfun(@(x) col(x(:,rleaves)),omegas,'UniformOutput',false);
omegas = [omegas{:}];
M = size(omegas,1);

% plot sampling pattern if desired
% figure; plot(reshape(omegas(:,1),[],length(rleaves)),reshape(omegas(:,2),[],length(rleaves)));
% drawnow;

% use exact DTFT of high-res image to get data points
kData = dtft(im_ref.*(N^2/Nx/Ny),omegas.*repmat(N./[Nx,Ny],M,1),[Nx,Ny]./2,Nx*Ny > 2^13)./sqrt(N*N);

% add some noise to it
SNR = 40; % dB
sigma = 10^(-SNR/20)*norm(kData(:))/sqrt(M*sum(abs(svd(noisecov))));
noisecov = noisecov.*sigma^2;
noise = sqrtm(noisecov./2)*complex(randn([NCha,M]),randn([NCha,M]));
kData = kData + noise.';
kData = reshape(kData,[],length(rleaves),NCha);

%% non-SPIRiT-related preprocessing

% compute density weighting
w_u = mri_density_comp(omegas.*(N/(2*pi)),'voronoi');
w_max_vals = (mod(rleaves([2:end,1]) - rleaves([end,1:end-1]),interleaves))/2;
%
% w_u = min(w_u,2*accel); % threshold weights
w_u = reshape(w_u,[],length(rleaves));
w_u = min(w_u,repmat(w_max_vals,[size(w_u,1),1]));

% create weighted and unweighted NUFFT operators
k_u = complex(omegas(:,1),omegas(:,2))./(2*pi); % in cycles instead of radians
k_u = reshape(k_u,[],length(rleaves));
NUFFTop_u = NUFFT(k_u, ones(size(w_u)), [0,0], [N,N]);
NUFFTop_uw = NUFFT(k_u, w_u, [0,0], [N,N]);

%% compute circulant preconditioner for NUFFT (takes a long time)
NUFFTgram_cs_spectrum = DTFT_circ_precon(omegas,[N,N]);

%% get ACS region for calibration

[omegas_cal{1},omegas_cal{2}] = ndgrid(((0:CalibSize(1)-1)-CalibSize(1)/2),(0:CalibSize(2)-1)-CalibSize(2)/2);
omegas_cal = [col(omegas_cal{1}),col(omegas_cal{2})].*(2*pi/N);

kCalib = dtft(im_ref.*(N^2/Nx/Ny),omegas_cal.*repmat(N./[Nx,Ny],size(omegas_cal,1),1),[Nx,Ny]./2,Nx*Ny > 2^13)./sqrt(N*N);
noise = sqrtm(noisecov./2)*complex(randn([NCha,size(omegas_cal,1)]),randn([NCha,size(omegas_cal,1)]));
kCalib = kCalib + noise.';
kCalib = reshape(kCalib,[CalibSize,NCha]);

%% perform SPIRiT calibration

kSize = [7,7];                  % SPIRiT Kernel size
CalibTyk = 0.02;                % Tykhonov regularization for calibration 0.01-0.05 recommended

kernel = zeros([kSize,NCha,NCha]);
[AtA,] = corrMatrix(kCalib,kSize);
fprintf(1,'Calibrating coil ');
cwidth = length(sprintf('%d',NCha));
for n=1:NCha
    fprintf(1,'%*d/%-*d...',cwidth,n,cwidth,NCha);
	kernel(:,:,:,n) = calibrate(AtA,kSize,NCha,n,CalibTyk);
    fprintf(1,repmat('\b',[1,2*cwidth+4]));
end
SPIRiTop = SPIRiT(kernel,'image',[N,N]); % init the SPIRiT Operator
fprintf(1,'done.\n');

%% conjugate phase recon (used as initialization later)
im_cp = NUFFTop_uw'*(kData.*repmat(sqrt(w_u),[1,1,NCha]));
% figure; imshow([sos(im_cp),sos(im_cp-im)],[0,1],'Border','tight'); title('Conjugate phase w/ dc'); colorbar();

im_init = im_cp;

%% regularizing parameters
wavWeight = 0.001;
tvWeight = 0.001;
lambda = 1;
wavelet_args = {'Daubechies',4,4};
tv_isotropic = true;

%% other parameters
if full_convergence
    ADMMiters = 2000; % longer for full convergence
else
    ADMMiters = 125; % initial results
end
PCGiters = 4;

%% perform many iterations of preconditioned ADMM or SB to get "baseline" image

K_NUFFT = 4000; K_SPIRiT = 20; spfraction_max = 0.4; % close to "optimal"
mus = [];  % ADMM tuning parameters

[mus,cond_NUFFT,cond_SPIRiT,spfraction] = L1SPIRiT_ADMM_select_parameter(K_NUFFT,K_SPIRiT,spfraction_max,NUFFTgram_cs_spectrum,im_init,SPIRiTop,lambda,wavWeight,tvWeight,wavelet_args,tv_isotropic,mus);

switch upper(baseline_algorithm)
    case 'SB'
        baseline_iters = 20000; % seems plenty

        ticker(60);
        [im_baseline,info_baseline] = L1SPIRiT_SB_NonCartesian(kData,im_init,SPIRiTop,NUFFTop_u,lambda,wavWeight,tvWeight,mus(2),'output',1,'maxiters',baseline_iters,'miniters',baseline_iters,'update_x_cg_opts',{'maxiters',PCGiters,'miniters',PCGiters},'update_x_precon','diag','wavelet_args',wavelet_args,'tv_isotropic',tv_isotropic,'test_convergence',true,'evalfun',nrmsefun);
        im_baseline = im_baseline(1).x; % baseline is final image
        info_baseline = rmfield(info_baseline,'x_cg_info'); % don't need to keep this around
        ticker('reset');

    case 'ADMM'
        baseline_iters = 4*ADMMiters; % a lot faster than SB

        ticker(60);
        [im_baseline,info_baseline] = L1SPIRiT_ADMM_NonCartesian(kData,im_init,SPIRiTop,NUFFTop_u,lambda,wavWeight,tvWeight,mus,'output',1,'maxiters',baseline_iters,'miniters',baseline_iters,'update_x_cg_opts',{'maxiters',PCGiters,'miniters',PCGiters},'update_x_precon',NUFFTgram_cs_spectrum,'update_z_solve','exact','wavelet_args',wavelet_args,'tv_isotropic',tv_isotropic,'test_convergence',true,'evalfun',nrmsefun);
        im_baseline = im_baseline(1).x; % baseline is final image
        info_baseline = rmfield(info_baseline,'x_cg_info'); % don't need to keep this around
        ticker('reset');
    otherwise
        error('Unknown baseline algorithm.');
end

nrmsdfun = @(x) norm(col(x - im_baseline))/norm(col(im_baseline));

%% examine "baseline" image

im_baseline_sqr = sos(im_baseline);

im_baseline_err_sqr = sos(im_baseline - im);

fprintf(1,'NRMSE difference of baseline: %g.\n',nrmsefun(im_baseline)); drawnow;

%% do ADMM with preconditioning and iterative z-update for a range of mu's
if mus_test
    matlabpool('open','local');

    mu1_range = logspace(-2,0,11);
    mu2_range = logspace(-3,-1,11);

    [mu1_grid,mu2_grid] = ndgrid(mu1_range,mu2_range);
    info_admm_precon_ziter_mus = cell(numel(mu1_grid),1);

    parfor imus = 1:numel(mu1_grid)
        SPIRiTop_use = SPIRiT(kernel,'image',[N,N]);
        NUFFTop_use = NUFFT(k_u, ones(size(w_u)), [0,0], [N,N]);
        [~,info_admm_precon_ziter_mus{imus}] = L1SPIRiT_ADMM_NonCartesian(kData,im_init,SPIRiTop_use,NUFFTop_use,lambda,wavWeight,tvWeight,[mu1_grid(imus);mu2_grid(imus)],'output',1,'maxiters',125,'miniters',125,'update_x_cg_opts',{'maxiters',PCGiters,'miniters',PCGiters},'update_x_precon',NUFFTgram_cs_spectrum,'update_z_solve','cg','update_z_cg_opts',{'maxiters',PCGiters,'miniters',PCGiters},'update_z_precon','diag','wavelet_args',wavelet_args,'tv_isotropic',tv_isotropic,'evalfun',nrmsdfun,'test_convergence',true);
    end

    info_admm_precon_ziter_mus = reshape(info_admm_precon_ziter_mus,[length(mu1_range),length(mu2_range)]);

    matlabpool('close');

    %% Brainweb contour plots (Fig. 4 for image 1)
    nrmsds_mus = cellfun(@(x) x.evaloutputs(end),info_admm_precon_ziter_mus);
    [nrmsd_min,nrmsd_minind] = min(nrmsds_mus(:));
    mu1_min = mu1_grid(nrmsd_minind);
    mu2_min = mu2_grid(nrmsd_minind);
    figure; [~,hc] = contour(mu1_grid,mu2_grid,20*log10(nrmsds_mus)); hold on;
    set(hc,'ShowText','on');
    plot(mu1_min,mu2_min,'kx');
    xlabel('\mu_z'); ylabel('\mu_W'); set(gca,'XScale','log','YScale','log');

    fmin = min(info_baseline.objs); % from baseline
    objs_mus = cellfun(@(x) min(x.objs)-fmin,info_admm_precon_ziter_mus);
    [obj_min,obj_minind] = min(objs_mus(:));
    mu1_min = mu1_grid(obj_minind);
    mu2_min = mu2_grid(obj_minind);
    figure; [~,hc] = contour(mu1_grid,mu2_grid,log10(objs_mus)); hold on;
    set(hc,'ShowText','on');
    plot(mu1_min,mu2_min,'kx');
    xlabel('\mu_z'); ylabel('\mu_W'); set(gca,'XScale','log','YScale','log');

end

%% do NLCG-based recon
maxiters = ADMMiters*PCGiters;    % number of reconstruction iterations

[im_nlcg,info_nlcg] = L1SPIRiT_NLCG_NonCartesian(kData,im_init,SPIRiTop,NUFFTop_u,lambda,wavWeight,tvWeight,'tsave',tStore,'output',0,'maxiters',maxiters,'miniters',maxiters,'wavelet_args',wavelet_args,'tv_isotropic',tv_isotropic,'evalfun',nrmsdfun);
[~,ind] = min(abs([im_nlcg(:).t] - tStore));
im_nlcg = im_nlcg(ind).x;
im_nlcg_err = im_nlcg - im_baseline;
im_nlcg_sqr = sos(im_nlcg);
im_nlcg_err_sqr = sos(im_nlcg_err);

%% do Split Bregman recon without preconditioning
[im_sb,info_sb] = L1SPIRiT_SB_NonCartesian(kData,im_init,SPIRiTop,NUFFTop_u,lambda,wavWeight,tvWeight,mus(2),'tsave',tStore,'output',0,'maxiters',ADMMiters,'miniters',ADMMiters,'update_x_cg_opts',{'maxiters',PCGiters,'miniters',PCGiters},'update_x_precon','','wavelet_args',wavelet_args,'tv_isotropic',tv_isotropic,'evalfun',nrmsdfun);
[~,ind] = min(abs([im_sb(:).t] - tStore));
im_sb = im_sb(ind).x;
im_sb_err = im_sb - im_baseline;
im_sb_sqr = sos(im_sb);
im_sb_err_sqr = sos(im_sb_err);

%% do Split Bregman recon with diagonal preconditioning
[im_sb_precon,info_sb_precon] = L1SPIRiT_SB_NonCartesian(kData,im_init,SPIRiTop,NUFFTop_u,lambda,wavWeight,tvWeight,mus(2),'tsave',tStore,'output',0,'maxiters',ADMMiters,'miniters',ADMMiters,'update_x_cg_opts',{'maxiters',PCGiters,'miniters',PCGiters},'update_x_precon','diag','wavelet_args',wavelet_args,'tv_isotropic',tv_isotropic,'evalfun',nrmsdfun);
[~,ind] = min(abs([im_sb_precon(:).t] - tStore));
im_sb_precon = im_sb_precon(ind).x;
im_sb_precon_err = im_sb_precon - im_baseline;
im_sb_precon_sqr = sos(im_sb_precon);
im_sb_precon_err_sqr = sos(im_sb_precon_err);

%% do ADMM-based recon without preconditioning but with exact z-updates
[im_admm,info_admm] = L1SPIRiT_ADMM_NonCartesian(kData,im_init,SPIRiTop,NUFFTop_u,lambda,wavWeight,tvWeight,mus,'tsave',tStore,'output',0,'maxiters',ADMMiters,'miniters',ADMMiters,'update_x_cg_opts',{'maxiters',PCGiters,'miniters',PCGiters},'update_x_precon',[],'update_z_solve','exact','wavelet_args',wavelet_args,'tv_isotropic',tv_isotropic,'evalfun',nrmsdfun);
[~,ind] = min(abs([im_admm(:).t] - tStore));
im_admm = im_admm(ind).x;
im_admm_err = im_admm - im_baseline;
im_admm_sqr = sos(im_admm);
im_admm_err_sqr = sos(im_admm_err);

%% do ADMM-based recon with preconditioning and exact z-update
[im_admm_precon_zexact,info_admm_precon_zexact] = L1SPIRiT_ADMM_NonCartesian(kData,im_init,SPIRiTop,NUFFTop_u,lambda,wavWeight,tvWeight,mus,'tsave',tStore,'output',0,'maxiters',ADMMiters,'miniters',ADMMiters,'update_x_cg_opts',{'maxiters',PCGiters,'miniters',PCGiters},'update_x_precon',NUFFTgram_cs_spectrum,'update_z_solve','exact','wavelet_args',wavelet_args,'tv_isotropic',tv_isotropic,'evalfun',nrmsdfun);
[~,ind] = min(abs([im_admm_precon_zexact(:).t] - tStore));
im_admm_precon_zexact = im_admm_precon_zexact(ind).x;
im_admm_precon_zexact_err = im_admm_precon_zexact - im_baseline;
im_admm_precon_zexact_sqr = sos(im_admm_precon_zexact);
im_admm_precon_zexact_err_sqr = sos(im_admm_precon_zexact_err);

%% do ADMM-based recon with preconditioning and iterative z-update
[im_admm_precon_ziter,info_admm_precon_ziter] = L1SPIRiT_ADMM_NonCartesian(kData,im_init,SPIRiTop,NUFFTop_u,lambda,wavWeight,tvWeight,mus,'tsave',tStore,'output',0,'maxiters',ADMMiters,'miniters',ADMMiters,'update_x_cg_opts',{'maxiters',PCGiters,'miniters',PCGiters},'update_x_precon',NUFFTgram_cs_spectrum,'update_z_solve','cg','update_z_cg_opts',{'maxiters',PCGiters,'miniters',PCGiters},'update_z_precon','diag','wavelet_args',wavelet_args,'tv_isotropic',tv_isotropic,'evalfun',nrmsdfun);
[~,ind] = min(abs([im_admm_precon_ziter(:).t] - tStore));
im_admm_precon_ziter = im_admm_precon_ziter(ind).x;
im_admm_precon_ziter_err = im_admm_precon_ziter - im_baseline;
im_admm_precon_ziter_sqr = sos(im_admm_precon_ziter);
im_admm_precon_ziter_err_sqr = sos(im_admm_precon_ziter_err);

%% reference image (Fig. 1 in both paper and supplement)
im_sqr = flipimfun(im_sqr);
figure; subplot('Position',[0,0,0.5,1.0]); imshow(im_sqr,window_image,'Border','tight'); hold on;
hrect = rectangle('Position',irect([2,1,4,3]),'LineWidth',2,'EdgeColor','w');
subplot('Position',[0.5,0,0.5,1.0]); imshow(im_sqr(irect(1) + (0:irect(3)-1),irect(2) + (0:irect(4)-1)),window_image,'Border','tight','InitialMagnification',mag); %xlabel('(a)');

%% recons, difference images (Figs. 6 and 7 in supplement)
window_diff = window_image./window_diff_scale;

im_nlcg_sqr = flipimfun(im_nlcg_sqr);
im_diff = flipimfun(sos(im_nlcg - im));
figure; subplot('Position',[0,0,0.5,1.0]); imshow(im_nlcg_sqr,window_image,'Border','tight'); hold on;
hin = axes('Position',[0.25,0,0.25,0.5]);
imagesc(im_nlcg_sqr(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,'XColor','w','YColor','w');
subplot('Position',[0.5,0,0.5,1.0]); imshow(im_diff,window_diff,'Border','tight'); hold on;
text(N,1,sprintf('%d\\times ',window_diff_scale),'Color','w','FontSize',20,'HorizontalAlignment','right','LineStyle','none','VerticalAlignment','top');
hin = axes('Position',[0.75,0,0.25,0.5]);
imagesc(im_diff(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,'XColor','w','YColor','w');

im_sb_precon_sqr = flipimfun(im_sb_precon_sqr);
im_diff = flipimfun(sos(im_sb_precon - im));
figure; subplot('Position',[0,0,0.5,1.0]); imshow(im_sb_precon_sqr,window_image,'Border','tight'); hold on;
hin = axes('Position',[0.25,0,0.25,0.5]);
imagesc(im_sb_precon_sqr(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,'XColor','w','YColor','w');
subplot('Position',[0.5,0,0.5,1.0]); imshow(im_diff,window_diff,'Border','tight'); hold on;
text(N,1,sprintf('%d\\times ',window_diff_scale),'Color','w','FontSize',20,'HorizontalAlignment','right','LineStyle','none','VerticalAlignment','top');
hin = axes('Position',[0.75,0,0.25,0.5]);
imagesc(im_diff(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,'XColor','w','YColor','w');

im_admm_precon_ziter_sqr = flipimfun(im_admm_precon_ziter_sqr);
im_diff = flipimfun(sos(im_admm_precon_ziter - im));
figure; subplot('Position',[0,0,0.5,1.0]); imshow(im_admm_precon_ziter_sqr,window_image,'Border','tight'); hold on;
hin = axes('Position',[0.25,0,0.25,0.5]);
imagesc(im_admm_precon_ziter_sqr(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,'XColor','w','YColor','w');
subplot('Position',[0.5,0,0.5,1.0]); imshow(im_diff,window_diff,'Border','tight'); hold on;
text(N,1,sprintf('%d\\times ',window_diff_scale),'Color','w','FontSize',20,'HorizontalAlignment','right','LineStyle','none','VerticalAlignment','top');
hin = axes('Position',[0.75,0,0.25,0.5]);
imagesc(im_diff(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,'XColor','w','YColor','w');

im_admm_precon_zexact_sqr = flipimfun(im_admm_precon_zexact_sqr);
im_diff = flipimfun(sos(im_admm_precon_zexact - im));
figure; subplot('Position',[0,0,0.5,1.0]); imshow(im_admm_precon_zexact_sqr,window_image,'Border','tight'); hold on;
hin = axes('Position',[0.25,0,0.25,0.5]);
imagesc(im_admm_precon_zexact_sqr(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,'XColor','w','YColor','w');
subplot('Position',[0.5,0,0.5,1.0]); imshow(im_diff,window_diff,'Border','tight'); hold on;
text(N,1,sprintf('%d\\times ',window_diff_scale),'Color','w','FontSize',20,'HorizontalAlignment','right','LineStyle','none','VerticalAlignment','top');
hin = axes('Position',[0.75,0,0.25,0.5]);
imagesc(im_diff(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,'XColor','w','YColor','w');

%% convergence plots

% for full convergence plots, in Figs. 4-5 in supplement

fmin = min(info_baseline.objs); % from baseline

figure;
if full_convergence
    subplot(2,1,1);
    plot([0;cumsum(info_sb.times)],20*log10(info_sb.evaloutputs(:)),[0;cumsum(info_sb_precon.times)],20*log10(info_sb_precon.evaloutputs(:)),[0;cumsum(info_nlcg.times)],20*log10(info_nlcg.evaloutputs(:)),[0;cumsum(info_admm.times)],20*log10(info_admm.evaloutputs(:)),[0;cumsum(info_admm_precon_ziter.times)],20*log10(info_admm_precon_ziter.evaloutputs(:)),[0;cumsum(info_admm_precon_zexact.times)],20*log10(info_admm_precon_zexact.evaloutputs(:)));
    xlim([0,min([sum(info_sb.times),sum(info_sb_precon.times),sum(info_nlcg.times),sum(info_admm.times),sum(info_admm_precon_ziter.times),sum(info_admm_precon_zexact.times)])]); % show it all
    switch imfile
        case 1
            ylim([-200,50]); set(gca,'YTick',-200:50:50,'Box','off');
        case 2
            ylim([-200,50]); set(gca,'YTick',-200:50:50,'Box','off');
    end
    xlabel('Time (s)'); ylabel('NRMSDs (dB)');
    subplot(2,1,2);
end

hs = semilogy([0;cumsum(info_sb.times)],max(0,info_sb.objs(:)-fmin),[0;cumsum(info_sb_precon.times)],max(0,info_sb_precon.objs(:)-fmin),[0;cumsum(info_nlcg.times)],max(0,info_nlcg.fs(:)-fmin),[0;cumsum(info_admm.times)],max(0,info_admm.objs(:)-fmin),[0;cumsum(info_admm_precon_ziter.times)],max(0,info_admm_precon_ziter.objs(:)-fmin),[0;cumsum(info_admm_precon_zexact.times)],max(0,info_admm_precon_zexact.objs(:)-fmin));
if full_convergence
    xlim([0,min([sum(info_sb.times),sum(info_sb_precon.times),sum(info_nlcg.times),sum(info_admm.times),sum(info_admm_precon_ziter.times),sum(info_admm_precon_zexact.times)])]); % show it all
    switch imfile
        case 1
            ylim([1e-15,1e6]); set(gca,'YTick',logspace(-15,5,5),'Box','off');
        case 2
            ylim([1e-15,1e5]); set(gca,'YTick',logspace(-15,5,5),'Box','off');
    end
else
    hold on;
    [~,ind] = min(abs([0;cumsum(info_sb.times)] - tStore)); hs(1) = semilogy(sum(info_sb.times(1:ind-1)),info_sb.objs(ind)-fmin,'-*','Color',get(hs(1),'Color'));
    [~,ind] = min(abs([0;cumsum(info_sb_precon.times)] - tStore)); hs(2) = semilogy(sum(info_sb_precon.times(1:ind-1)),info_sb_precon.objs(ind)-fmin,'-+','Color',get(hs(2),'Color'));
    [~,ind] = min(abs([0;cumsum(info_nlcg.times)] - tStore)); hs(3) = semilogy(sum(info_nlcg.times(1:ind-1)),info_nlcg.fs(ind)-fmin,'-o','Color',get(hs(3),'Color'));
    [~,ind] = min(abs([0;cumsum(info_admm.times)] - tStore)); hs(4) = semilogy(sum(info_admm.times(1:ind-1)),info_admm.objs(ind)-fmin,'-x','Color',get(hs(4),'Color'),'MarkerSize',8);
    [~,ind] = min(abs([0;cumsum(info_admm_precon_ziter.times)] - tStore)); hs(5) = semilogy(sum(info_admm_precon_ziter.times(1:ind-1)),info_admm_precon_ziter.objs(ind)-fmin,'-s','Color',get(hs(5),'Color'));
    [~,ind] = min(abs([0;cumsum(info_admm_precon_zexact.times)] - tStore)); hs(6) = semilogy(sum(info_admm_precon_zexact.times(1:ind-1)),info_admm_precon_zexact.objs(ind)-fmin,'-d','Color',get(hs(6),'Color'));
    xlim([0,tStore*1.2]); % show some times
    switch imfile
        case 1
            ylim([1e-6,1e4]); set(gca,'YTick',logspace(-6,4,6),'Box','off');
        case 2
            ylim([1e-6,1e4]); set(gca,'YTick',logspace(-6,4,6),'Box','off');
    end
end
xlabel('Time (s)'); ylabel('f({\bfX}) - f^{opt}');
legend(hs([1,4,2,3,5,6]),{'SB','ADMM','SB + precon','NLCG','ADMM + precon (z-cg)','ADMM + precon (z-exact)'},'Location','SouthWest');
