% we implement irls iterative solve for regulrizing Ax=b
% via minimization of J(x) = 2*\sum_{k=1}^N \lambda_k |x_k|^p + norm(A*xn-b,2)^l;
% we have A,x,b; norm(A)<1
% Sergey Voronin, see cg based acceleration for inverse problems
function [x_sol,num_iters_taken,tElapsed,residuals,Js,xnx_errors] = irls_sys_solve(A,AtA,b,x,l,p,lambda,x0,TOL,maxiters) 

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

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

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

fprintf('in irls_solve, starting loop..\n');
fprintf('maxiters = %d, lambda = %f\n', maxiters, lambda);
for i=1:maxiters
    if quit_loop ~= 1
            if mod(i,10) == 0
                fprintf('doing irls iter %d of %d with lambda = %f\n', i, maxiters, lambda);
            end

            num_iters_taken = i;

            % compute residual
            residuals(i) = norm(A*xn - b,2)^2;
            Js(i) = 2*lambda*norm(xn,1) + norm(A*xn-b,2)^2;
            xnx_errors(i) = 100*norm(xn-x)/norm(x);

            % record old solution
            xn1 = xn;

            % compute weights
            for ind=1:n
                power = (2-p)/2;
                wns(ind) = 1./(1 + (xn(ind)^2 + en^2)^power);
            end

            % construct D
            D = eye(n,n);
            for k=1:n
                D(k,k) = sqrt(0.5*lambda*p*wns(k));
            end

            % construct R
            R = eye(m,m);
            res = A*xn - b;
            eps = 1e-4;
            for j=1:m
                rj = abs(res(j));
                if rj < eps
                    rj = eps;
                end  
                R(j,j) = l*rj^(l-2);
            end

            %xn = lsqr(A'*(R*A) + D*D, A'*(R*b));
            xn = pcg(A'*(R*A) + D*D, A'*(R*b));


            % update en
            %en = min( en, sqrt(norm(xn - xn1) + gamma^n) );
            %en = min( en, sqrt(norm(xn - xn1) + gamma^n) );
            %en = en/1.1;
            en = norm(xn - xn1);

            % 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_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);

