% This code generates results in TABLE 1 of the paper
%--------------------------------------------------------------------------
% [1]  Joy A, Paul J S. "Multi-channel MR Image Reconstruction using
%      Statistically Optimized Non-Linear Diffusion", Magnetic Resonance in 
%      Medicine, June 2017, doi:10.1002/mrm.26774.
% Email: ajin.joy@iiitmk.ac.in, j.paul@iiitmk.ac.in
% 
% Coded by: Ajin Joy on March 2017
%
% These codes and data are shared to help other researchers understand
% and reproduce our work.
% 
% This implimentation needs PANO and PBDW codes avaliable for download at
% http://www.quxiaobo.org/ (not included)
%--------------------------------------------------------------------------
clc; clear all; close all;
disp('Time to run this code:');
disp('~15 minutes on a quad core(2.4GHz) machine running MATLAB 2015a'); 
params.gamma  = .1;  % moderate diffusion strength. (0<gamma<.25]
params.c      = 1;   % maximum consistency. [0<c<1]
params.numitr = 100; % maximum number of iterations

load datasets;
I             = dataset{1}.I;
[OutPM,OutTV,PM_PBDW,TV_PBDW,GT_PBDW,PM_PANO,TV_PANO,GT_PANO,...
    PMtime,TVtime,PM2time,TV2time,GT2time,PM3time,TV3time,GT3time] = deal(cell(1,3));
RLNE          =@(x,y) norm(x(:)-y(:),2)./norm(x(:),2);

% COMPUTE RESULTS
%--------------------------------------------------------------------------
for i=1:3
    if i==1,     mask = dataset{1}.mask_x3;
    elseif i==2, mask = dataset{1}.mask_x4;
    elseif i==3, mask = dataset{1}.mask_x5;
    end
    fprintf('        UNDERSAMPLING x%d\n',i+2);
    % PM and TV reconstruction to provide reference images for PBDW and PANO
    disp('::::generating reference images')
    disp(' ----running PM diffusion');
    params.gamma = .1; 
    DF           =@(dI,alpha_sq) 1./(1 + (dI.*conj(dI))/(alpha_sq));tic;
    OutPM{i}     = NLdiffusionMRI_stripped(I,mask,DF,params);
    PMtime{i}    = toc;
    disp('  Completed.')
    disp(' ----running TV recon');
    DF           =@(dI,dummy) 1./((dI.*conj(dI) + 1e-10).^(1/2));tic;
    params.gamma = searchgamma(I,mask,DF,params,50,50);% 50 searches max(50 iterations per search)
    OutTV{i}     = NLdiffusionMRI_stripped(I,mask,DF,params);
    TVtime{i}    = toc;
    disp('  Completed.')
    %----------------------------------------------------------------------
    % Learned reconstruction using PBDW
    disp('::::Learned recon using PBDW');
    disp(' ----reference image used: PM (~1-2minutes)');
    tic;PBDW     = PBDWsol(OutPM{i}.Irecon); % PBDW operator
    PM_PBDW{i}   = learnedrecon(I,mask,PBDW,OutPM{i}.Irecon);
    PM2time{i}   = toc;
    disp('  Completed.')
    disp(' ----reference image used: TV (~1-2minutes)');
    tic;PBDW     = PBDWsol(OutTV{i}.Irecon); % PBDW operator
    TV_PBDW{i}   = learnedrecon(I,mask,PBDW,OutTV{i}.Irecon);
    TV2time{i}   = toc;
    disp('  Completed.')
    disp(' ----reference image used: GT (~1-2minutes)');
    tic;PBDW     = PBDWsol(I); % PBDW operator
    GT_PBDW{i}   = learnedrecon(I,mask,PBDW,I);
    GT2time{i}   = toc;
    disp('  Completed.')
    %----------------------------------------------------------------------
    % Learned reconstruction using PANO
    disp('::::Learned recon using PANO');
    disp(' ----reference image used: PM');tic;
    PM_PANO{i}   = PANOsol(I,mask,OutPM{i}.Irecon);
    PM3time{i}   = toc;
    fprintf(repmat('\b',1,271));
    disp('  Completed.')
    disp(' ----reference image used: TV');tic;
    TV_PANO{i}   = PANOsol(I,mask,OutTV{i}.Irecon);
    TV3time{i}   = toc;
    fprintf(repmat('\b',1,271));
    disp('  Completed.')
    disp(' ----reference image used: GT (ground truth)');tic;
    GT_PANO{i}   = PANOsol(I,mask,I);
    GT3time{i}   = toc;
    fprintf(repmat('\b',1,271));
    disp('  Completed.')
    %----------------------------------------------------------------------
end

for i=1:3
    % DISPLAY RESULTS
    %----------------------------------------------------------------------
    fprintf('\n::::RESULTS for undersampling x%d\n',i+2);
    fprintf(' [PM]      RLNE: %.4f Elapsed time taken: %.2f seconds\n',RLNE(abs(I),abs(OutPM{i}.Irecon)),PMtime{i});
    fprintf(' [TV]      RLNE: %.4f Elapsed time taken: %.2f seconds\n',RLNE(abs(I),abs(OutTV{i}.Irecon)),TVtime{i});
    fprintf('\n');
    fprintf(' [PM_PBDW] RLNE: %.4f Elapsed time taken: %.2f seconds\n',RLNE(abs(I),abs(PM_PBDW{i}.Irecon)),PM2time{i});
    fprintf(' [TV_PBDW] RLNE: %.4f Elapsed time taken: %.2f seconds\n',RLNE(abs(I),abs(TV_PBDW{i}.Irecon)),TV2time{i});
    fprintf(' [GT_PBDW] RLNE: %.4f Elapsed time taken: %.2f seconds\n',RLNE(abs(I),abs(GT_PBDW{i}.Irecon)),GT2time{i});
    fprintf('\n');
    fprintf(' [PM_PANO] RLNE: %.4f Elapsed time taken: %.2f seconds\n',RLNE(abs(I),abs(PM_PANO{i}.Irecon)),PM3time{i});
    fprintf(' [TV_PANO] RLNE: %.4f Elapsed time taken: %.2f seconds\n',RLNE(abs(I),abs(TV_PANO{i}.Irecon)),TV3time{i});
    fprintf(' [GT_PANO] RLNE: %.4f Elapsed time taken: %.2f seconds\n',RLNE(abs(I),abs(GT_PANO{i}.Irecon)),GT3time{i});
    %----------------------------------------------------------------------


%     % DISPLAY FIGURES
%     %----------------------------------------------------------------------
%     cl1   = cat(1,abs(OutPM{i}.Irecon),abs(PM_PBDW{i}.Irecon),abs(PM_PANO{i}.Irecon));
%     cl2   = cat(1,abs(OutTV{i}.Irecon),abs(PM_PBDW{i}.Irecon),abs(PM_PANO{i}.Irecon));
%     cl3   = cat(1,abs(I),abs(GT_PBDW{i}.Irecon),abs(GT_PANO{i}.Irecon));
%     rwcl1 = cat(2,cl1,cl2,cl3);
%     figure;imagesc(rwcl1);colormap gray;axis image;axis off;drawnow
%     title(sprintf('COLUMNS: PM, TV and GroundTruth Images--- ROWS: Ref Image, PBDWrecon, PANOrecon-- undersampling x%d',i+2),'fontsize',12);
%     suptitle('Supporting figure: PBDW and PANO reconstruction from PM diffusion and TV');
%     cl1   = cat(1,abs(I)-abs(OutPM{i}.Irecon),abs(I)-abs(PM_PBDW{i}.Irecon),abs(I)-abs(PM_PANO{i}.Irecon));
%     cl2   = cat(1,abs(I)-abs(OutTV{i}.Irecon),abs(I)-abs(PM_PBDW{i}.Irecon),abs(I)-abs(PM_PANO{i}.Irecon));
%     cl3   = cat(1,abs(I)-abs(I),abs(I)-abs(GT_PBDW{i}.Irecon),abs(I)-abs(GT_PANO{i}.Irecon));
%     rwcl2 = cat(2,cl1,cl2,cl3);
%     figure;imagesc(rwcl2);colormap hsv;axis image;axis off;colorbar;drawnow
%     title(sprintf('COLUMNS: PM, TV and GroundTruth Images--- ROWS: Ref Image, PBDWrecon, PANOrecon-- undersampling x%d',i+2),'fontsize',12);
%     suptitle('Supporting figure: Error images');
%     %----------------------------------------------------------------------
end