function [Out] = NLdiffusionMRI(I,mask,DF,params)
% Input parameters:
%   I     : coil image
%   mask  : sampling mask (same size of I)
%   DF    : diffusivity function
%           for PM diffusion use DF =@(dI,alpha_sq) 1./(1 + (dI.*conj(dI))/(alpha_sq));
%           for TV diffusion use DF =@(dI,dummy) 1./((dI.*conj(dI) + 1e-10).^(1/2));
% Output parameters:
%   Out.Irecon : Reconstructed image
%   Out.Izf    : Fourier transform reconstruction(zero filled image)
%   Out.err    : RLNE measure
% 
%--------------------------------------------------------------------------
% [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
%--------------------------------------------------------------------------
if nargin<3, DF =@(dI,alpha_sq) 1./(1 + (dI.*conj(dI))/(alpha_sq)); end
if nargin<4, params=[]; end

if ~isfield(params,'numitr'), numitr = 100;  else, numitr = params.numitr; end
if ~isfield(params,'gamma'),  gamma  = .1;   else, gamma  = params.gamma;  end
if ~isfield(params,'c'),      c      = 1;    else, c      = params.c;      end
if ~isfield(params,'scale'),  scale  = .5^2; else, scale  = params.scale;  end

[M,N]   = size(I);
RLNE    =@(x,y) norm(x(:)-y(:),2)./norm(x(:),2);
FTf     =@(x)   (1/sqrt(M*N)*fftshift(fft2(ifftshift(x)))).*mask; % fourier transform with sampling
FTi     =@(x)   sqrt(M*N)*fftshift(ifft2(ifftshift(x.*mask))); % inverse fourier transform
K       = FTf(I);
Izf     = FTi(K);
Irec    = Izf;
err     = zeros(1,numitr);

%-------------------------------------------------------------
figure(99),subplot(1,2,1);imagesc(abs(Izf));axis image
colormap gray;title('Input');drawnow
%-------------------------------------------------------------

k=1; t_last=1; Irec_last=Izf;
while (k<=numitr)
    % biasing ----------------------------------------------------
    update1   = FTi(c*(FTf(Irec_last)-K));
    Irec      = Irec - update1;
    
    % diffusion denoising ----------------------------------------
    dI        = FDf(Irec);
    alpha_sq  = scale*mad(dI(:))^2;	% estimating alpha
    update2   = gamma*(FDi(dI.*DF(dI,alpha_sq)));
    Irec      = Irec - update2;	

    % acceleration using FISTA -----------------------------------
    Irec2     = Irec;   
    t_now     = (1 + sqrt(1 + 4*t_last*t_last))/2;
    Irec      = Irec + (t_last - 1)/t_now*(Irec - Irec_last);
    t_last    = t_now;
    Irec_last = Irec2;
    
    [err(k)]  = RLNE(abs(I),abs(Irec)); % error measure

    %-------------------------------------------------------------
    figure(99),subplot(1,2,2);imagesc(abs(Irec));axis image;colormap gray;
    title(sprintf('Reconstructing: Itr# %d/%d',k,numitr));drawnow
    %-------------------------------------------------------------
    k=k+1;
end
close figure 99
% output
Out.Irecon = Irec;
Out.Izf    = Izf;
Out.err    = err(1:k);


%------------------- calling function --------------------
function Out = FDf(x)
% derivatives
Out(:,:,1) = x([2:end,end],:) - x;
Out(:,:,2) = x(:,[2:end,end]) - x;

%------------------- calling function --------------------
function Out = FDi(x)
% divergence
dx          = x(:,:,1);
dy          = x(:,:,2);
Outx        = dx([1,1:end-1],:) - dx;
Outy        = dy(:,[1,1:end-1]) - dy;
Outx(1,:)   = -dx(1,:);
Outy(:,1)   = -dy(:,1);
Outx(end,:) = dx(end-1,:);
Outy(:,end) = dy(:,end-1);
Out         = Outx + Outy;
