function [imgE,energyi,ISNRi,energyo,ISNRo,Lambda,discr]=tvmm(img_noisy,h,varargin)
% function   [imgE]=tvmm_debluring(img_noisy,h,...
%            'optional_parameter_name1',value1,'optional_parameter_name2', value2,...);
%
% Adaptive Total Variation-based image deconvolution with
% a majorization-minimization approach.
%  
% Written by: Joao Oliveira
% email: joao.oliveira@lx.it.pt
% SITE: www.lx.it.pt/~jpaos
% Date: 21/12/2007
%  
%        
% ========================== INPUT PARAMETERS (required) =================
% Parameter     Values
% name          and description
% ========================================================================
% img_noisy		(double) Noisy blured image of size ny. 
% h             (double) Blur kernel.
% lambda		Regularization parameter (which is multiplied by the
%               TV penalty). In case of adaptative lambda, set this
%               variable as empty (ex: []).
%
% ======================== OPTIONAL INPUT PARAMETERS ====================
% Parameter     Values
% name          and description
% =======================================================================
% MM_iter      (double) The number of outer loop iterations. 
%               Default: 10 
% MM_thrld
%
% S_iter
%
% S_thrld
%
%
%
% cg_iter       (double) The max number of inner CG iterations.
%               Default: 100
% cg_thrld      (double) Conjugate Gradient threshold.
%               Default: 1e-5;
% image         (double) Original image.
% displayIm     ({'yes','no'}) if 'yes' is passed  the restored imaged
%               is displayed along the CG iterations
%               Default: 'no'
% info_energyi  ({'yes','no'}) if 'yes' is passed the isnr and energy of
%               the objective function is displayed inside the inner loop.
%               Default: 'no'
% info_energyo  ({'yes','no'}) if 'yes' is passed the energy of
%               the objective function is displayed at each outter loop.
%               Default: 'no'
% info_ISNRi    ({'yes','no'}) if 'yes' is passed the isnr is displayed.
%               Default: 'no'
%               Requires: image
% info_ISNRo    ({'yes','no'}) if 'yes' is passed the isnr is displayed.
%               Default: 'no'
%               Requires: image
% x_0           (double) initial image iteration.
%               Default: Wiener filter
% sigma         (double) Variance of noise. If not set, Donoho MAD's rule
%               will be used instead, in case of adaptative lambda.
%
%       ===================================================================
%       The following functions can be provided in order to overwrite the 
%       internal ones. Recall that, by default, all calculations are
%       performed with circular convolutions. (see Technical Report)
%       ===================================================================
%
% mult_H        (function_handle) Function handle to the function that 
%               performes H*x. This function must accept two parameters:
%                       x : image to apply the convolution kernel h
%                       h : Blur kernel
% mult_Ht       (function_handle) Function handle to the function that 
%               performes Ht*x (Ht = H transpose). This function must 
%               accept two parameters:
%                       x : image to apply the convolution kernel h
%                       h : Blur kernel
% diffh         (function_handle) Function handle to the function that 
%               computes the horizontal differences of an image x.
% diffv         (function_handle) Function handle to the function that 
%               computes the vertical differences of an image x.
% diffht         (function_handle) Function handle to the function that 
%               computes the horizontal differences (transposed) of an 
%               image x.
% diffh         (function_handle) Function handle to the function that 
%               computes the vertical differences (transposed) of an 
%               image x.
%
% alpha         Gamma parameter
%               Default: 0.5
%
% beta          Gamma parameter
%               Default: 1e5
%
% ====================== OUTPUT PARAMETERS ==============================
% imgE              Estimated image
%
% ==================== OPTIONAL OUTPUT PARAMETERS =======================
% Parameter     Values
% name          and description
%
% Note: The parameters appear in the following order, according to the
%       required input arguments flags
% =======================================================================
% energyi           Energy of inner loop iterations.
%                   Required input arguments: 'info_energyi' and 'image'.
% ISNRi             Improvement Signal-to-noise ratio of inner loop iterations.
%                   Required input arguments: 'info_ISNRi' and 'image'.
% energyo           Energy of outer loop iterations.
%                   Required input arguments: 'info_energyo' and 'image'.
% ISNRo             Improvement Signal-to-noise ratio of outer loop iterations.
%                   Required input arguments: 'info_ISNRo' and 'image'.
% Lambda            Lambda estimated/used at each MM iteration.
% ============= EXAMPLES =========================================



% *******************************************************
% *********************** M U D A R *********************
% *******************************************************



%
% Normal adaptive use:
% [img_estimated]=tvmm(img_noisy,h,sigma)
% 
% Normal use (lambda given):
% [img_estimated]=tvmm(img_noisy,h,'lambda',lambda)
% 
% Display restored images along CG iterations:
% [img_estimated]=tvmm(img_noisy,h,lambda,sigma,...
%                       'displayIm','yes')
%
% Perform ISNR calculations of the outer loop iterations:
% [img_estimated,energyi,ISNRi,energyo,ISNRo]=tvmm(img_noisy,h,...
%                       h,lambda,sigma,'info_SNRo','yes','image',image);
%
% 

% ======================= DEFAULT PARAMETERS ===================
global mH mHt diffh diffv diffht diffvt
MM_iter = 10;          % Number of outer loop iterations
MM_thrld = 1e-5;
S_iter = 10;
S_thrld = 1e-5;
CG_iter = 200;          % Number of inner CG iterations
CG_thrld = 1e-5;        % CG threshold
displayIm = 0;
image=[];
sigma=[];
lambda=[];
Lambda=[];
info_energyi='no';
info_energyo='no';
info_ISNRi='no';
info_ISNRo='no';
info_int = 0;
info_ext = 0;
energyi=0;
ISNRi=0;
energyo=0;
ISNRo=0;
mH=@conv2c;
mHt=@conv2c;
diffh=@f_diffh;
diffv=@f_diffv;
diffht=@f_diffht;
diffvt=@f_diffvt;
x_0=[];
alpha=0.5;
theta=0.5;
beta=1;
% Quotient of the number of outter loop iterations vs update of lambda 
% parameter
% Default: 1 (i.e., it actualizes the lambda in each outter loop iteration)



% ====================== INPUT PARAMETERS ========================
% Test for number of required parametres
if (nargin-length(varargin)) ~= 2
     error('Wrong number of required parameters');
end




if (iscell(varargin{1}))
    varargin=varargin{1};
%     nargin=length(varargin);
end
% Read the optional parameters
if (rem(length(varargin),2)==1)
    error('Optional parameters should always go by pairs');
else
    fprintf('Reading optional parameters:\n');
    for i=1:2:(length(varargin)-1)
        % change the value of parameter
        switch varargin{i}
            case 'MM_iter'                     % Outer loop iterations
                MM_iter = varargin{i+1};
                fprintf('  MM iter iterations = %g\n',MM_iter);
            case 'MM_thrld'                     % Outer loop iterations
                MM_thrld = varargin{i+1};
                fprintf('  MM thrld iterations = %g\n',MM_thrld);
            case 'S_iter'                     % Outer loop iterations
                S_iter = varargin{i+1};
                fprintf('  MM outer iterations = %g\n',MM_iter);
            case 'S_thrld'                     % Outer loop iterations
                S_thrld = varargin{i+1};
                fprintf('  MM outer iterations = %g\n',S_thrld);
            case 'CG_iter'                     % Outer loop iterations
                CG_iter = varargin{i+1};
                fprintf('  MM outer iterations = %g\n',CG_iter);
            case 'CG_thrld'                     % Outer loop iterations
                CG_thrld = varargin{i+1};
                fprintf('  MM outer iterations = %g\n',CG_thrld);
            case 'displayIm'                    % display sucessive xe estimates
                if (isequal(varargin{i+1},'yes'))
                    displayIm=1;
                    fprintf('  Display image = ok\n',varargin{i+1});
                end
            case 'image'                        % Original image
                image = varargin{i+1};
                fprintf('  Original imaged supplied = yes\n');
            case 'info_ISNRi'                   % display PSNR inside inner loop
                info_ISNRi = varargin{i+1};
                fprintf('  Inner loop ISNR calculations = %s\n',varargin{i+1});
            case 'info_ISNRo'                   % display PSNR on outer loop
                info_ISNRo = varargin{i+1};
                fprintf('  Outer loop ISNR calculations = %s\n',varargin{i+1});
            case 'info_energyi'                 % display energy inside inner loop
                info_energyi = varargin{i+1};
                fprintf('  Inner loop ISNR calculations = %s\n',varargin{i+1});
            case 'info_energyo'                 % display energy on outer loop
                info_energyo = varargin{i+1};
                fprintf('  Inner loop ISNR calculations = %s\n',varargin{i+1});
            case 'x_0'                          % Initial image iteration
                x_0 = varargin{i+1};
                fprintf('  Initial iteration (x_0) supplied = yes\n');
            case 'mult_H'                       % Function handle to perform H*x
                mH = varargin{i+1};
                fprintf('  H operator supplied = yes\n');
            case 'mult_Ht'                      % Function handle to perform Ht*x
                mHt = varargin{i+1};
                fprintf(' Ht operator supplied = yes\n');
            case 'diffh'                   % External function to perform 
                diffh = varargin{i+1};     % horizontal differences    
                fprintf('  Diff horizontal operator supplied = ok\n');
            case 'diffv'                   % External function to perform 
                diffv = varargin{i+1};     % vertical differences    
                fprintf('  Diff vertical operator supplied = ok\n');
            case 'diffht'                  % External function to perform 
                diffht = varargin{i+1};    % horizontal differences (transposed)   
                fprintf('  Diff horizontal transpose operator supplied = ok\n');
            case 'diffvt'                  % External function to perform 
                diffvt = varargin{i+1};    % vertical differences (transposed)  
                fprintf('  Diff vertical transpose operator supplied\n');
            case 'sigma'                   % Noise Variance
                sigma = varargin{i+1};
                fprintf('  sigma = %g\n',varargin{i+1});
            case 'lambda'                   % Lambda parameter
                lambda = varargin{i+1};
                fprintf('  lambda = %g\n',varargin{i+1});
            case 'alpha'
                alpha = varargin{i+1};
                fprintf('  alpha = %g\n',varargin{i+1});
            case 'beta'
                beta = varargin{i+1};
                fprintf('  beta = %g\n',varargin{i+1});
            case 'theta'
                theta = varargin{i+1};
                fprintf('  theta = %g\n',varargin{i+1});
            case 'num_act'
                num_act = varargin{i+1};
                fprintf('  Lambda: ratio btw inner/outer actualization = %g\n',varargin{i+1});
            otherwise
                % Hmmm, something wrong with the parameter string
                error(['Unrecognized parameter: ''' varargin{i} '''']);
        end;
    end;
    fprintf('Done.\n');
end
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

% Parameters pos-validations (after load optional parameters)

% Adaptative lambda
if isempty(lambda)
    auto_lambda=1;
    fprintf('Adaptative lambda simulation!\n');
else
    auto_lambda=0;
end

if (isequal(info_ISNRi,'yes')||isequal(info_ISNRo,'yes'))&&isempty(image)
    error('ISNR calculations: original image not found');
end

% Which internal info to display...
if isequal(info_energyi ,'yes')
    info_int=info_int+1;
end
if isequal(info_ISNRi,'yes')
    info_int=info_int+2;
end

% Which external info to display...
if isequal(info_energyo ,'yes')
    info_ext=info_ext+1;
end
if isequal(info_ISNRo,'yes')
    info_ext=info_ext+2;
end

% Initializations
[hy,hx]=size(h);
[imy,imx]=size(img_noisy);

% Initialization with Wiener filter

h_aux=h;
h_aux(imy,imx)=0;
h_aux=imshift(h_aux,-floor(hy/2),-floor(hx/2));
db = fft2(h_aux);               % eigenvalues of B
varx = var(img_noisy(:));       % approximate var of img_noisy
aux = diffh(img_noisy);

if isempty(sigma)
    sigma =  median(abs(aux(:)-median(aux(:))))/.6745/2;
    fprintf('Noise variance not given! MAD Rule: sigma = %g\n',sigma);
end

if (isempty(x_0))
    x_old = real(fft2(db./(db.^2+10*sigma^2/varx).*ifft2(img_noisy)));
else
    x_old=x_0;
end


% Initial Improve Signal-to-noise Ratio
if bitand(info_ext,2)==2 || bitand(info_int,2)==2
    ISNRi=[20*log10(norm(img_noisy-image,'fro')/norm(x_old-image,'fro'))];
    ISNRo=ISNRi;
end

% Initial energy
if bitand(info_int,1)==1 || bitand(info_ext,1)==1
    energyi=[norm(img_noisy-mH(x_old,h),'fro')^2+lambda*sum(sum(sqrt(diffh(x_old).^2+diffv(x_old).^2)))];
    energyo=energyi;
end

% Initializing gamma distributions
% fprintf('First lambda = %g \n', 2*alpha*imy*imx*sigma.^2/(sum(sum(sqrt(diffh(x_old).^2+diffv(x_old).^2)))+beta));
%==================== TV MM ======================
i=1;
ti=clock;
wij=diffh(x_old).^2+diffv(x_old).^2;

Wij=1./(sqrt(wij));
if (auto_lambda)
    lambda=2*(alpha + theta*imy*imx)*sigma.^2/(sum(sum(sqrt(wij)))+beta);
end
lambda_old=lambda;
decr_MM=2*MM_thrld;
while(i<=MM_iter && decr_MM>MM_thrld )

    if (info_ext>0)
        fprintf('MM %g of %g with new lambda=%g\n',i,MM_iter,lambda);
    end
    decr_S=2*S_thrld;
    s=1;
    while (s<=S_iter && decr_S>S_thrld)
        t1=clock;
        if (info_ext>0)
            fprintf('MM %g of %g \n',s,S_iter);
        end

        Wij=1./(sqrt(wij));
        [x_new,n_it_cg,cost_i,isnri]=cg_tv(Wij,img_noisy,lambda,x_old,h,CG_iter,CG_thrld,displayIm,info_int,image,alpha,beta,sigma);
        wij=diffh(x_new).^2+diffv(x_new).^2;
        t2=clock;
        % Keep track of internal data
        if bitand(info_int,1)==1
            energyi=[energyi cost_i];
        end
        if bitand(info_int,2)==2
            ISNRi=[ISNRi isnri];
        end

        % Display external simulation info
        if (info_ext>0)

            fprintf('MM %g of %g: CG iters = %g',i,MM_iter,n_it_cg);
            %             fprintf(', num = %g ',length(indices));
            fprintf(', lambda = %g',lambda);
            if bitand(info_ext,1)==1
                %energyo=[energyo norm(img_noisy-mH(x_new,h),'fro')^2+lambda*sum(sum(sqrt(wij)))];
                energyo=[energyo norm(img_noisy-mH(x_new,h),'fro')^2+2*(alpha+theta*numel(img_noisy))*sigma^2*log(sum(sum(sqrt(diffh(x_new).^2+diffv(x_new).^2)))+beta)];
                fprintf(', energy = %g',energyo(end));
            end
            if bitand(info_ext,2)==2
                ISNRo=[ISNRo 10*log10(norm(img_noisy-image,'fro').^2/norm(x_new-image,'fro').^2)];
                fprintf(', ISNR = %g dB',ISNRo(end));
            end
            fprintf(', time = %g secs\n',etime(t2,t1));
        end
        Lambda=[Lambda, lambda];
        if s>2
            decr_S=abs(norm(x_new-x_old,'fro')/norm(x_old,'fro'));
        end
        x_old=x_new;
        s=s+1;
    end

    Wij=1./(sqrt(wij));
    if (auto_lambda)
        lambda=2*(alpha + theta*imy*imx)*sigma.^2/(sum(sum(sqrt(wij)))+beta);
    end

    if i>2
        decr_MM=abs(norm(lambda-lambda_old,'fro')/norm(lambda_old,'fro'));
    end

    lambda_old=lambda;
    i=i+1;
end


tf=clock;
if info_ext>0 || info_int>0
    fprintf('Simulation time = %g secs \n',etime(tf,ti));
end

imgE=x_new;
discr=norm(img_noisy-conv2c(imgE,h),'fro')^2/numel(imgE);

% =============    Auxiliary functions   =================

function sol=f_diffh(x)
h=[0 1 -1];
sol=conv2c(x,h);


function sol=f_diffht(x)
h=[-1 1];
sol=conv2c(x,h);


function sol=f_diffv(x)
h=[0 1 -1]';
sol=conv2c(x,h);


function sol=f_diffvt(x)
h=[-1 1]';
sol=conv2c(x,h);

function y = conv2c(x,h)
% Circular 2D convolution
x=wraparound(x,h);
y=conv2(x,h,'valid');


function y = wraparound(x, m)
% Extend x so as to wrap around on both axes, sufficient to allow a
% "valid" convolution with m to return the cyclical convolution.
% We assume mask origin near centre of mask for compatibility with
% "same" option.
[mx, nx] = size(x);
[mm, nm] = size(m);
if mm > mx | nm > nx
    error('Mask does not fit inside array')
end

mo = floor((1+mm)/2); no = floor((1+nm)/2);  % reflected mask origin
ml = mo-1;            nl = no-1;             % mask left/above origin
mr = mm-mo;           nr = nm-no;            % mask right/below origin
me = mx-ml+1;         ne = nx-nl+1;          % reflected margin in input
mt = mx+ml;           nt = nx+nl;            % top of image in output
my = mx+mm-1;         ny = nx+nm-1;          % output size

y = zeros(my, ny);
y(mo:mt, no:nt) = x;      % central region
if ml > 0
    y(1:ml, no:nt) = x(me:mx, :);                   % top side
    if nl > 0
        y(1:ml, 1:nl) = x(me:mx, ne:nx);            % top left corner
    end
    if nr > 0
        y(1:ml, nt+1:ny) = x(me:mx, 1:nr);          % top right corner
    end
end
if mr > 0
    y(mt+1:my, no:nt) = x(1:mr, :);                 % bottom side
    if nl > 0
        y(mt+1:my, 1:nl) = x(1:mr, ne:nx);          % bottom left corner
    end
    if nr > 0
        y(mt+1:my, nt+1:ny) = x(1:mr, 1:nr);        % bottom right corner
    end
end
if nl > 0
    y(mo:mt, 1:nl) = x(:, ne:nx);                   % left side
end
if nr > 0
    y(mo:mt, nt+1:ny) = x(:, 1:nr);                 % right side
end




function [x_new,i,iter_energy,isnri]=cg_tv(Wij,b,lambda,x_old,h,imax,cg_thrld,displayIm,info_int,image,alpha,beta,sigma)
% Function cg_tv applies the conjugate gradient
% method in order to solve the linear system Ax=b, where
% A = (D'*V*D*lamda/2 +H'H)
%    D - toeplitz matrix
%    V - Diagonal matrix
%
% Inputs parameters are:
%    Wij : 
%    lambda  : lambda parameter
%    b  : particular solution b
%    x_old : starting value for x
%    h : blur kernel
%    imax : maximum CG iterations
%    cg_thrld : threshold for stoping CG
%    displayIm : display restored images along CG iterations
%    info_int : perform (or not) calculations for ISNR or energy
%    image : original image for ISNR calculations
% Output parameters are:
%    x_new : solution of conjudate gradient
%    i : number of iterations used
%    energy : energy function of all iterations 
%    isnri : Improve Signal-to-noise ratio of all iterations
%
% Written by: Joao Oliveira
% email: joao.oliveira@lx.it.pt
% Date: 28/09/2005

global mH mHt diffh diffv diffht diffvt
% Inicialization
i=0;
iter_energy=[];
isnri=[];
decr=2*cg_thrld;


% Initial Conjugate gradient parameters
% System: Az=y, y=Ht*x

% y=conv2c(b,h);
y=mHt(b,h);
% Multiply x_old by [lambda/2*D'*V*D + Ht*H]
r=y-(lambda/2)*(diffht(Wij.*diffh(x_old))+diffvt(Wij.*diffv(x_old)))-mHt(mH(x_old,h),h);
d=r;
delta_new=r(:)'*r(:);
delta0=delta_new;

while(i<imax && decr>cg_thrld)
% while(i<imax)
    
    i=i+1;
    % q=A*d; A=[lambda/2*D'*V*D + Ht*H]
    q=(lambda/2)*(diffht(Wij.*diffh(d))+diffvt(Wij.*diffv(d)))+mHt(mH(d,h),h);
    alfa=delta_new/(d(:)'*q(:));

    x_new=x_old+alfa*d;

    % Display internal simulation info
    if (info_int>0)
        fprintf('\t\tCG iteration %g',i);
        if bitand(info_int,1)==1
            %iter_energy=[iter_energy norm(b-mH(x_new,h),'fro')^2+lambda*sum(sum(sqrt(diffh(x_new).^2+diffv(x_new).^2)))];
            iter_energy=[iter_energy norm(b-mH(x_new,h),'fro')^2+2*(alpha+theta*numel(b))*sigma^2*log(sum(sum(sqrt(diffh(x_new).^2+diffv(x_new).^2)))+beta)];
            fprintf(', energy = %g',iter_energy(end));
        end
        if bitand(info_int,2)==2
            isnri=[isnri 20*log10(norm(b-image,'fro')/norm(x_new-image,'fro'))];
            fprintf(', ISNR = %g dB',isnri(end));
        end
        fprintf('\n');
    end

    % Display restored image
    if (displayIm)
        figure(7);imagesc(x_new);colormap gray; drawnow;
%         figure(8);imagesc(x_new-image); colormap gray; drawnow;
    end
    
    r=r-alfa*q;
    delta_old=delta_new;
    delta_new=r(:)'*r(:);
    beta=delta_new/delta_old;
    d=r+beta*d;

    if i>2
        decr=abs(norm(x_new-x_old,'fro')/norm(x_old,'fro'));
    end
    x_old=x_new;
end


function  z = imshift(x,shx,shy)
%         z = imshift(x,shx,shy);
%
% circularly shift image by nx and ny 
%
% ----------- Input parameters -----------------
% x		    -> original 2D signal 
% shx		-> x-shift 
% shy		-> y-shift 
%
% ----------- Output parameters -----------------
% z 		-> Shifted image
   
[ny,nx] = size(x);
I=sqrt(-1);

% shift dft 
fs=exp(-I*2*pi/ny*shy*(0:(ny-1))')*exp(-I*2*pi/nx*shx*(0:(nx-1)));

% symmetric shift
z=real(ifft2(fft2(x).*fs));

 
function [z,i]=cgtv2(Wij,w,lambda,z,imax)
% Solve Az=Dw
% Inicialization
global diffh diffv diffht diffvt mH mHt
i=1;
[dimy,dimx]=size(z);


% Initial Conjugate gradient parameters
% System: Az=y
Dw=[diffh(w); diffv(w)];

Dt=diffht(z(1:(dimy/2),:))+diffvt(z(dimy/2+1:end,:));
r=Dw-[diffh(Dt); diffv(Dt)]-2/lambda*[Wij.*z(1:(dimy/2),:); Wij.*z(dimy/2+1:end,:)];
d=r;
delta_new=r(:)'*r(:);

% Iterations

while(i<=imax)
   
    % q=A*d;
    Dt=diffht(d(1:(dimy/2),:))+diffvt(d(dimy/2+1:end,:));
    q=[diffh(Dt); diffv(Dt)]+2/lambda*[Wij.*d(1:(dimy/2),:); Wij.*d(dimy/2+1:end,:)];
    alfa=delta_new/(d(:)'*q(:));
    z=z+alfa*d;
    r=r-alfa*q;
    delta_old=delta_new;
    delta_new=r(:)'*r(:);
    beta=delta_new/delta_old;
    d=r+beta*d;
    i=i+1;
end





