% 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 = {'/y/wellerds/MRI_data/L1SPIRiT_ADMM_v1.0/spgr_spiral.mat'; ...
    '/y/wellerds/MRI_data/L1SPIRiT_ADMM_v1.0/ssfp_radial.mat'};
imfile = 1; % 8-channel SPGR spiral data
% imfile = 2; % 8-channel SSFP radial data

full_convergence = false; % set to true to get supplementary convergence plots; false gets plots in paper

% path setup
if ~exist('SPIRiT', 'class')
%	addpath('../SPIRiT_v0.1/');
	addpath('../ESPIRiT/');
end
if ~exist('zpad.m', 'file')
%	addpath('../SPIRiT_v0.1/utils/');
	addpath('../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 = 25;

        irect = [109,81,52,44];
        window_image = [0,0.5];
        window_diff_scale = 5;
        
        flipimfun = @(im) im(:,end:-1:1).';
    case 2
        accel = 4;
        CalibSize = [30,30];            % size of the calibration region
        tStore = 25;
        
        irect = [33,109,56,48];
        window_image = [0,0.5];
        window_diff_scale = 5;
        
        flipimfun = @(im) im.';
%         accel = 4;
%         CalibSize = [36,36];            % size of the calibration region
%         tStore = 100; % in seconds
    otherwise
        error('unknown file selected.');
end

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

% basic initialization
rng('default');

%% load real data with coil sensitivities
load(imfiles{imfile},'data','im','k','w');
im = double(im);
data = double(data);
k = double(k);
w = double(w);

[Nx,Ny,NCha] = size(im);
im_sqr = sos(im);

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 = size(k,2);
idx = sort(randperm(interleaves,ceil(interleaves/accel))); % random undersampling

% undersample by selecting interleaves
k_u = k(:,idx);
w_reweighting = (mod(idx([2:end,1]) - idx([end,1:end-1]),interleaves))/2;
w_u = w(:,idx).*repmat(w_reweighting,[size(w,1),1]);
data = data(:,idx,:);
omegas = [real(k_u(:)),imag(k_u(:))].*(2*pi);
M = numel(k_u);

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

NUFFTop_u = NUFFT(k_u, ones(size(k_u)), [0,0], [Nx,Ny]);
NUFFTop_uw = NUFFT(k_u, w_u, [0,0], [Nx,Ny]);

data = data./sqrt(prod([Nx,Ny])); % rescale data to match NUFFT object

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

%% get ACS region for calibration

kCalib = fft2c(im); % from reference image
kCalib = crop(kCalib,[CalibSize,NCha]); % crop center k-space for calibration

%% 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',[Nx,Ny]); % init the SPIRiT Operator
fprintf(1,'done.\n');

%% conjugate phase recon (used as initialization later)
im_cp = NUFFTop_uw'*(data.*repmat(sqrt(w_u),[1,1,NCha]));

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(data,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(data,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 NLCG-based recon
maxiters = ADMMiters*PCGiters;    % number of reconstruction iterations

[im_nlcg,info_nlcg] = L1SPIRiT_NLCG_NonCartesian(data,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(data,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(data,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(data,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(data,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(data,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. 7 and 8 in paper)
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(Nx,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(Nx,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(Nx,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(Nx,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 regular convergence plots, in Figs. 5 and 6 in paper
% for full convergence plots, in Figs. 2-3 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([-150,50]); set(gca,'YTick',-150:50:50,'Box','off');
        case 2
            ylim([-150,50]); set(gca,'YTick',-150: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-11,1e6]); set(gca,'YTick',logspace(-10,5,4),'Box','off');
        case 2
            ylim([1e-15,1e6]); 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*2.4]); % show some times
    switch imfile
        case 1
            ylim([1e-6,1e5]); set(gca,'YTick',logspace(-6,4,6),'Box','off');
        case 2
            ylim([1e-5,1e4]); set(gca,'YTick',logspace(-4,4,5),'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');
