function output = TVCMRI(input,varargin)
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% Authors: Shiqian Ma, Wotao Yin
% Date : 10/11/2007
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

% minimize alpha*TV(Phi'*x) + beta*||x||_1 + 0.5*||Ax-b||_2^2

t0 = cputime;
% Read the optional parameters and Initialization
[x,y,xs,...     % initial x and y
 crit_x,crit_y,f,...
 n1,n2,N,... % dimension
 alpha,beta,alpha_final,beta_final,eta,...
 A,b,Phi,Atb,...
 tau,TV_iso,TV_eps,...
 xtol,ytol,gtol,ftol,maxitr,output,print_lvl] = initilization(input,varargin{:}); t1 = cputime;

TV_iso=1;v1=0;g1=0;v2=0;v3=0;r3=0;g3=0;iitr=0;f_prev=inf;crit_x_prev=inf;

get_grad(); save_stats();
% alpha = 1e-3; beta = 3.5e-2; 
alpha = input.alpha; beta = input.beta; 

iterno=1; 
output.funv(iterno)=f; output.xrate(iterno)=1;
output.snr(iterno)=0;  output.xtime(iterno)=0;
maxitr=input.no;
% iterno=0; 
oitr = 1;           % outer iter counter
for itr = 1:maxitr  % total iter counter
    iterno=iterno+1;
    iitr = iitr+1;    % inner iter counter
    xp = x; yp = y;
    xh = x - tau*(alpha*(Phi*(Lt(y)))+g3);
    x = shrink1D(xh,zeros(N,1),tau*beta);
    xIm=Phi'*x;
    yh = y + tau*L(xIm); y = shrink2D(yh,1/tau);    
    get_grad(); save_stats();
    output.funv(iterno)=f;
    output.xrate(iterno)=norm(x-input.xs)/norm(input.xs);
    output.snr(iterno)=snr(Phi'*x, input.f);
    output.xtime(iterno)=cputime-t0;
%     if print_lvl>=2; fprintf(1,' total=%d inner=%d',itr,iitr); print_stats(); end
    
%     if meet_inner_crit()
    if itr>=input.no
        if beta == beta_final; 
            t2 = cputime; save_output(); return;
        else
            alpha = max(eta*alpha,alpha_final); output.alpha = [output.alpha alpha];
            beta = max(eta*beta,beta_final); output.beta = [output.beta beta];
            oitr = oitr + 1; iitr = 0;
        end
    end

end

t2 = cputime;
save_output(); 
return;

%% %%%% nested helper functions %%%%
    function [x,y,xs,...     % initial x and y
              crit_x,crit_y,f,...
              n1,n2,N,... % dimension
              alpha,beta,alpha_final,beta_final,eta,...
              A,b,Phi,Atb,...
              tau,TV_iso,TV_eps,...
              xtol,ytol,gtol,ftol,maxitr,output,print_lvl] = initilization(input,varargin)
        xs  = [];
        n1  = input.n1;
        n2  = input.n2;
        
        alpha_final = input.alpha;
        beta_final  = input.beta;
        N   = n1*n2;
        A   = input.A;
        Phi = input.Phi;
        b   = input.b;
        Atb = A'*b;
        % default options
        eta = 1/4000;  % the rate of reducing alpha and beta
        TV_iso = false;
        TV_eps = 0;
        tau = 0.8;
        xtol = 1e-5;
        ytol = 1e-5;
        gtol = inf;
        ftol = 1e-5;
        maxitr = 100;
        print_lvl = 1;
        % read options
        if (rem(length(varargin),2)==1) ; error('Opts go by pairs');
        else for ii=1:2:(length(varargin)-1) 
            switch varargin{ii}
            case 'true x'
                xs = varargin{ii+1}(:);
            case 'print level'
                print_lvl = varargin{ii+1}(:);
            otherwise
                error(['Unrecognized option: ''' varargin{ii} '''']);
            end
            end
        end
        % initialize x and y
        x = zeros(N,1);
        y = zeros(N,2);
        crit_x = inf;
        crit_y = inf;
        f = inf;
        % initialize alpha and beta
        beta = max(eta*norm(Atb,inf),beta_final);
        alpha = alpha_final*beta/beta_final;
        % initialize stats
        output.x_hist = [];
        output.y_hist = [];
        output.f_hist = [];
        output.pv1_hist = [];
        output.pv2_hist = [];
        output.pv3_hist = [];
        output.step_x_hist = [];
        output.crit_x  = [];
        output.crit_y  = [];
        output.crit_g  = [];
        output.alpha = [];
        output.beta = [];
    end

    function get_grad()
        v1 = alpha*TV(Phi'*x);
        g1 = alpha*(Phi*(Lt(y)));
        v2 = beta*norm(x,1);
        Ax = A*x;
        g3 = A'*Ax - Atb;
        r3 = Ax - b; 
        v3 = 0.5*(norm(r3(:)).^2);
        f_prev = f;
        f = v1 + v2 + v3;
    end

    function save_stats()
        
        %  output.x_hist = [output.x_hist x]; % too memory consuming
        %  output.y_hist = [output.y_hist y]; % too memory consuming
        output.f_hist = [output.f_hist f];
        output.pv1_hist = [output.pv1_hist v1/alpha];
        output.pv2_hist = [output.pv2_hist v2/beta];
        output.pv3_hist = [output.pv3_hist v3];
    end

    function print_stats()
            if ~isempty(xs); rel_err = norm(x-xs)/norm(xs); else rel_err=NaN; end
            fprintf(1,' obj=%3e nonzeros=%7d rel_err=%.3e',...
                        f, nnz(abs(x)>1e-4), rel_err);
    end

    function bStop = meet_inner_crit()
        step_x = norm(x-xp); crit_x_prev = crit_x; crit_x = step_x/max(1,norm(xp));
        step_y = norm(y(:)-yp(:)); crit_y = step_y/max(1,norm(yp(:)));
        crit_f = (f_prev - f)/max(1,f_prev);
        crit_g = norm(g1+g3,'inf') - beta;
        if print_lvl>=2; fprintf(1,' crit_x=%.3e crit_y=%.3e crit_g=%.3e\n',...
                                   crit_x,crit_y,crit_g); end
        bStop = ( crit_x<xtol... % *sqrt(beta/beta_final)...
               && crit_y<ytol... % *sqrt(alpha/alpha_final)...
               )...              % && crit_g < gtol)...
               || crit_f<ftol...%*sqrt(beta/beta_final)...   
               || crit_x >= crit_x_prev...
               || iitr >= input.maxiitr;
%           
      bStop =   iitr > input.maxiitr;
           
        output.step_x_hist = [output.step_x_hist step_x];
        output.crit_x = [output.crit_x crit_x];
        output.crit_y = [output.crit_y crit_y];
        output.crit_g = [output.crit_g crit_g];
    end

    function save_output()
        output.x = x;
        output.f = f;
        output.y = y;
        output.itr = itr;
        output.oitr = oitr;
        output.init_time = t1-t0;
        output.iter_time = t2-t1;
    end

    function p=Grad1(u)
        % backward finite difference along dim 1
        u = reshape(u,n1,n2);
        p = [u(1,:);diff(u,1,1);];
        p = p(:);
    end

    function pp=Grad1Star(u)
        % adjoint of backward finite difference along dim 1
        u = reshape(u,n1,n2);
        pp = [-diff(u,1,1);u(end,:)];
        pp = pp(:);
    end

    function q=Grad2(u)
        % backward finite difference along dim 2
        u = reshape(u,n1,n2);
        q = [u(:,1) diff(u,1,2)];
        q = q(:);
    end

    function qq=Grad2Star(u)
        % adjoint of the backward finite difference along dim 2
        u = reshape(u,n1,n2);
        qq = [-diff(u,1,2) u(:,end)];
        qq = qq(:);
    end

    function TVx=TV(x)
        % Total Variation norm of x, x is a n by n matrix
        grad_x = [Grad1(x) Grad2(x)];
        if TV_iso
            pt_sqsum = sum(grad_x.*grad_x,2);
            if TV_eps == 0; TVx = sum(sqrt(pt_sqsum)); else TVx = sum(sqrt(pt_sqsum+TV_eps)); end
        else
            TVx = norm(grad_x(:),1);
        end
    end

    function s1 = shrink1D(p,q,s)
        % shrink p toward q by up to s; p and q are vectors, s is scalar
        tmp = p - q; s1 = q + sign(tmp).*max(0,abs(tmp)-s);
    end

    function s2 = shrink2D(y,gamma)
        % y = [y1, y2], where y1 and y2 are N dimensional vectors
        % output s2 is N by 2 matrix
        V = sqrt(sum(y.*y,2));
        S = max(0,V-gamma);
        V(V==0) = 1; S = S./V;
        s2 = y - [S.*y(:,1) S.*y(:,2)];
    end

% function s2 = shrink2D(y,gamma)
%         % y = [y1, y2], where y1 and y2 are N dimensional vectors
%         % output s2 is N by 2 matrix
%         V = sqrt(sum(y.*y,2));
%         index2=find(V>gamma);
%         s2=y;
%         s2(index2, 1)=y(index2, 1)./V(index2)/gamma;
% end

    function y = L(x)
        y = [Grad1(x) Grad2(x)];  % y(pixel,coord)
    end

    function x = Lt(y)
        x = Grad1Star(y(:,1)) + Grad2Star(y(:,2));
    end

end
