% we implement thresholded FISTA with IRLS for regulrizing Ax=b
% via minimization of J(x) = 2*tau*norm(xn,1) + norm(A*xn-b,2)^2;
% we have A,x,b; norm(A)<1
% 
% here's how we do thresholded IRLS FISTA:
% 
% let T(x) = wn*( x + A'*b - A'*A*x )
% next,
% pick some x0 and iterate:
% x_{n+1}=T(x_n+\frac{t_{n}-1}{t_{n+1}} (x^{(n)}-x^{(n-1)}))
% where {t_n} is a sequence of numbers, generated by:
% t_{n+1}=\frac{1+\sqrt{1+4(t_{n})^2}}{2}$ and t_1=1
function [x_sol,num_iters_taken,tElapsed,residuals,Js,xnx_errors] = irls_fista_solve(A,b,x,tau,x0,TOL,maxiters) 

tStart = tic;
m = size(A,1);
n = size(A,2);

% generate sequence of tn's needed for thresholded FISTA
tns = repmat(0,maxiters+10,1);
tns(1) = 1;
for i=2:length(tns)
    tns(i) = (1 + sqrt(1 + 4*tns(i-1)^2))/2;
end

% calculate z = A^t*b
z = A'*b;

num_iters_taken = 0;
xn = x0;
xn1 = x0;
wns = ones(n,1);
en = 1e-1;
gamma = 0.1;

quit_loop = 0;
residuals = repmat(0,maxiters,1);
Js = repmat(0,maxiters,1);
xnx_errors = repmat(0,maxiters,1);


fprintf('in irls_fista_solve, starting loop..\n');
fprintf('maxiters = %d, tau = %f\n', maxiters, tau);
for i=1:maxiters
    if quit_loop ~= 1
        if mod(i,10) == 0
            fprintf('doing IRLS FISTA iter %d of %d with tau = %f\n', i, maxiters, tau);
        end
            %x_{n+1}=T(x_n+\frac{t_{n}-1}{t_{n+1}} (x^{(n)}-x^{(n-1)}));
            num_iters_taken = i;

            % compute residual
            residuals(i) = norm(A*xn - b,2)^2;
            Js(i) = 2*tau*norm(xn,1) + norm(A*xn-b,2)^2;
            xnx_errors(i) = 100*norm(xn-x)/norm(x);
            
            if i==1
                pn1 = x0;
                pn = xn;
            else
                pn1 = pn;
                pn = xn + ((tns(i)-1)/tns(i+1))*(xn - xn1);
            end
            
            xn1 = xn; % record x_{n-1}
            un1 = 2*tau*norm(xn1,1) + norm(A*xn1-b,2)^2;

            vn = A*pn;
            wn = A'*vn;
            %xn = softThreshold(pn + z - wn, tau);

            % compute weights
            for ind=1:length(wns)
                %wns(ind) = 1./max(abs(xn(ind)),en);
                wns(ind) = 1./(1 + tau/sqrt(pn(ind)^2 + en^2));
            end

            % compute new solution component wise 
            for ind=1:length(xn)
                xn(ind) = wns(ind) * ( pn(ind) + z(ind) - wn(ind) );
            end

            % update en
            en = norm(pn - pn1);

            % record solution
%            close all;
%            if i>=10
%                im_file = ['images/frame',num2str(i),'.jpg'];
%            else
%                im_file = ['images/frame0',num2str(i),'.jpg'];
%            end
%            figure(i);
%            imshow(reshape(xn,125,125));
%            print('-djpeg',im_file);
%            close all;
%    
%            % trim
%            cmd = ['convert ', im_file,' -fuzz 1% -trim +repage ', im_file];
%            system(cmd);

            % compute diff between un and un1 and quit if small
            if(i>3 && (100*(norm(xn - xn1)/norm(xn)))<TOL)
                fprintf('reached convergence..\n');
                quit_loop = 1;
            end
    else
        break;
    end
end
fprintf('done with irls_fista_solve.. num_iters_taken = %d\n', num_iters_taken); 

xnx_errors = xnx_errors(1:num_iters_taken);

% record solution
x_sol = xn;

% record time
tElapsed = toc(tStart);

