function [u, k, result, cond_MA, cond_M, cond_A, eig_MA ,eig_M, eig_A] = ...
   nma_CG(A,f,tol,max_iterations,precond,handlerAxes,opt)
%conjugate gradient with pre-conditioning solver

%file name: nma_CG.m
%
% This function solves Au=f using the method of
% conjugate gradient with pre-conditioning
%
% INPUT:
%
%  A: System matrix,  must be PSD matrix
%
%  f: RHS. If scalar 0, then inital guess u used to start the iteration
%     process will be a random vector.
%     If f is a vector, then if it is an all zero vector, then initial guess
%     for u will also be random, otherwise, initial guess u used will be
%     zero
%
%  tol: tolerance on the norm of the residual ||r|| (norm 2).
%       When the residual r=f-Au has its 2 norm become smaller than
%       this value, then the iterative process will stop and u is returned.
%
%  max_iterations: if this number is reached before tolerance is reached
%       then process is stopped and error is thrown.
%
%  precond: preconditioner to use. Valid values are
%    'NONE' no pre-conditioning is done
%    'SSOR' for symmetric SOR
%    'INCOMPLETE-CHOLESKY'  for incomplete Cholesky
%    'MULTI-GRID'   for multigrid
%  
%  h: handle to plotting axis. If not zero, then a plot is send
%     to this axes handler at each iteration to show current solution.
%
%
%  optinal 6th argument. If present, will represent the DROPTOL (the
%          drop tolerance) to use when selecting the incomplete Cholesky
%          for preconditioning. This input is ignored if method is not
%          'INCOMPLETE-CHOLESKY'. This is a non negative number
%          If method is 'INCOMPLETE-CHOLESKY' and this input is missing, then DROPTOL=0
%          is used, which produces the complete Cholesky factorization.
%
% Output:
%  u: approximate solution to Au=f
%  k: number of iterations used to reach the solution
%  result: result table. See below for layout of this table.
%  cond_MA : condition number of M^-1 * A
%  cond_M  : condition number of M
%  cond_A  : condition number of A
%  eig_MA  : vector of eigenvalues of M^-1*A
%  eig_M   : vector of eigenvalues of M
%  eig_A   : vector of eigenvalues of A
%
% EXAMPLE:
% see nma_TEST_CG.m for example on calling this function
%
% By Nasser M. Abbasi
% MATH 228A, UC Davis, Fall 2010
% December 2, 2010



%----------- START ARGUMENT CHECKING AND VALIDATION -------------------
if nargin ~= 6 && nargin ~= 7
   error('nma_CGP:: number of arguments must be 6 or 7');
end

if not(nma_ISSPD(A))
   error('nma_CGP:: A is not SPD');
end

if not(strcmpi(precond,'SSOR')) && not(strcmpi(precond,'INCOMPLETE-CHOLESKY'))...
      && not(strcmpi(precond,'MULTI-GRID')) && not(strcmpi(precond,'NONE'))
   error('nma_CGP:: method must be one of NONE, SSOR, INCOMPLETE-CHOLESKY or MULTI-GRID');
end

if nargin==7
   if opt<0
      error('nma_CGP:: droptol input must be > 0');
   else
      droptol=opt;
   end
else
   droptol=0;
end

if isscalar(f)
   if not(abs(f)<=eps)
      error('nma_CGP:: f not a scalar zero and not a vector');
   end
else
   [nRowF,nColF] = size(f(:));
   if nColF ~= 1
      error('nma_CGP:: f must be a vector');
   end
end

[N,~] = size(A);   % N is number of unknowns
if not(isscalar(f))
   if N ~= nRowF
      error(' nma_CGP:: f size not compatible with A size');
   end
end

if any(f)
   u = zeros(N,1);
else
   u = rand(N,1);
end

%----------- END ARGUMENT CHECKING AND VALIDATION -------------------

%--
%-- Allocate space for result table, and find mesh spacing for
%-- mesh norm calculation below
%-- result table have the following columns order
%--    k, |e|, ratio of |e| , |r|, ratio of |r|
%--
result = zeros(max_iterations,5); %
h = 1/(sqrt(size(A,1))+1);
nx = sqrt(size(A,1));

current_r = f-A*u;
omega     = 1;  % for smoother use 2*(1-pi*h), but 1 for our case

cond_MA=0;
cond_M=0;
cond_A=0;
eig_MA=0;
eig_M=0;
eig_A=0;

%[cond_MA,cond_M,cond_A,eig_MA ,eig_M, eig_A] = ...
%   find_condition_number(A,omega,precond,droptol);

current_z = solve_for_z(current_r,omega,A,precond,droptol);
current_p = current_z;
k         = 1;

while true
   
   result = update(result,k,u,current_r,h);
   
   old_r     = current_r;
   old_p     = current_p;
   old_z     = current_z;
   w         = A*old_p;
   alpha     = (old_z'*old_r)/(old_p'*w);
   u         = u + alpha*old_p;
   current_r = old_r - alpha * w;
   current_z = solve_for_z(current_r,omega,A,precond,droptol);
   beta      = (current_z' * current_r)/(old_z' * old_r);
   current_p = current_z + beta * old_p;
   
   %--
   %-- Check if we achived convergence. Make sure to use the mesh
   %-- norm 2, not just norm 2. For saftey, check against maximum
   %-- number of iteration. CG always converges, but this is good to
   %-- have in case of a bug in the code. An error is thrown if not
   %-- converged before hitting the maximum iterations
   %--
   if (sqrt(h)*norm(current_r,2)>tol) && k<=max_iterations
      k = k + 1;
   else
      result = update(result,k,u,current_r,h);
      break;
   end
   
   %-- uncomment below to see the convergence during running
   %-- use for animation later
   
   if handlerAxes ~=0
      mesh(handlerAxes,reshape(-u,nx,nx));
      title(sprintf('current solution, iteration %d',k));
      drawnow;
      %pause(.1);
   end
end

if k>max_iterations
   error('nma_CG_no_pre_condition_solver:: no convergence before max_iterations');
end

end

%-------------------------
% Function to fill in result table, called from main line above
% INPUT:
%    result: The table
%    k     : iteration number
%    u     : current solution at the iteration number
%    current_r : the residule at this iteration
% OUTPUT
%    result: The result table after filling it
%
function result = update(result,k,u,current_r,h)

result(k,1)= k;
result(k,2)= sqrt(h)*norm(u,2);
result(k,4)= sqrt(h)*norm(current_r,2);
if k > 1
   result(k,3) = result(k,2)/result(k-1,2);
   result(k,5) = result(k,4)/result(k-1,4);
end

end

%----------------------------------------------
% This function solves for Z in Mz=r
% INPUT:
%  r: the residual, the RHS
%  w: omega for SSOR, this should be 1
%  A: the system matrix
%  precond: Name of preconditioner to use
%  droptol: the srop tolerance to use when the preconditioner is
%           inclomplete cholesky
%
% OUTPUT:
% z: the solution to Mz=r
%
%
function z = solve_for_z(r,w,A,precond,droptol)

%precond='EXACT';  %for testing
%precond='MSSOR';  %for testing

N  = size(A,1);  % number of unknowns
nx = sqrt(N);    % internal grid points
n  = nx+2;
h  = 1/(nx+1);
rT = zeros(n);


if strcmpi(precond,'NONE')
   z = r;
   
elseif strcmpi(precond,'MULTI-GRID')
   METHOD = 5; % Gauss-Seidel red/black as preconditioner
   mu1    = 1; % number of pre-smooth
   mu2    = 1; % number of post-smooth
   rT(2:end-1,2:end-1) = reshape(r,nx,nx);
   z = nma_V_cycle(zeros(nx+2),rT,mu1,mu2,METHOD);
   z = reshape( z(2:end-1,2:end-1),N,1);
   
elseif strcmpi(precond,'INCOMPLETE-CHOLESKY')
   if droptol == 0
      R = chol(A);
   else
      R = cholinc(A,droptol);
   end
   
   z = R\(R'\r);
elseif strcmpi(precond,'SSOR')
   
   zk0 = zeros(n);
   zkh = zeros(n);
   zk1 = zeros(n);
   rT(2:end-1,2:end-1) = reshape(r,nx,nx);
   
   % forward SOR sweep
   for i=2:n-1
      for j=2:n-1
         zkh(i,j) = (w/4) * ( zkh(i-1,j) + zkh(i,j-1) + ...
            zk0(i+1,j) + zk0(i,j+1) - h^2*rT(i,j) ) ...
            + (1-w) * zk0(i,j);
      end
   end
   
   % reverse SOR sweep
   for i=n-1:-1:2
      for j=n-1:-1:2
         zk1(i,j) = (w/4) * (zkh(i-1,j)+zkh(i,j-1) + ...
            zk1(i+1,j)+zk1(i,j+1) - h^2*rT(i,j) ) + ...
            (1-w) * zkh(i,j);
      end
   end
   
   z = reshape( zk1(2:end-1,2:end-1),N,1);
   
elseif strcmpi(precond,'EXACT')
   d   = diag(diag(A));
   M   = (d-w*tril(A,-1)).*(d-w*triu(A,1));
   z=M\r;
elseif strcmpi(precond,'MSSOR')
   d=diag(diag(A));
   z = w*(2-w)* pinv(d+w*triu(A,1)) * d * pinv(d+w*tril(A,-1)) * r;
end

end

%-----------------------------------------
% Function to find condition numbers for A, M^-1*A
% and find the spectrum. This is for analysis only and not needed
% to actualy solve the problem
%
% INPUT:
%   A: system matrix
%   w: omega for SSOR, should be 1
%   droptol: the srop tolerance to use when the preconditioner is
%             inclomplete cholesky
%
% OUTPUT:
%   cond_MA: condition number for M**-1 * A
%   cond_M:  condition number for M
%   cond_A:  condition number for A
%   eig_MA:  all the eigenvalues for M**-1 * A
%   eig_M:   all the eigenvales for M
%   eig_A:   all the eigenvalues for A
%

function[cond_MA,cond_M,cond_A,eig_MA ,eig_M, eig_A] = ...
   find_condition_number(A,w,precond,droptol)

cond_MA = 0;
cond_M  = 0;
cond_A  = condest(A);
eig_A   = eig(A);
eig_M   = 0;
eig_MA  = 0;

if strcmpi(precond,'MULTI-GRID')
   %nx   = sqrt(size(A,1));
   %invM = find_inverse_M_for_multigrid(nx);
   cond_A  = condest(A);
   %cond_MA = condest( invM*A );
   
elseif strcmpi(precond,'INCOMPLETE-CHOLESKY')
   if droptol == 0
      M = chol(A);
   else
      M = cholinc(A,droptol);
   end
   
   
   cond_M  = condest(M);
   
   if(size(A,1)<16129)
      p = inv(M)*A;
      cond_MA = condest(p);
      eig_MA  = abs(eig(full(p)));
   else
      cond_MA = 0;
      eig_MA  = 0;
   end
   
elseif strcmpi(precond,'SSOR')
   d = diag(diag(A));
   M = (d + w*tril(A,-1) )*( d + w*triu(A,1) );
   if(size(A,1)<16129)
      p = inv(M)*A;
      cond_MA = condest(p);
      eig_MA  = abs(eig(full(p)));
   else
      cond_MA = 0;
      eig_MA  = 0;
   end
   cond_M  = condest(M);
   eig_M   = eig(M);
   
end

end

%----------------------------
% THis function is called to find M^-1 for the Multigrid case only
% in order to determine spectrum of M^-1 * A
% not needed to solve the CG problem, only used for
% analysis.
%
% INPUT:
%   n: number of grid points on one side of the grid
% OUTPUT:
%   M^-1
%
function invM = find_inverse_M_for_multigrid(n)

invM   = zeros(n^2);
METHOD = 5; % Gauss-Seidel red/black as preconditioner
mu1    = 1; % number of pre-smooth
mu2    = 1; % number of post-smooth

for i=1:n
   f    = zeros(n^2,1);
   f(i) = 1;
   f    = reshape(f,n,n);
   z = nma_V_cycle(zeros(n),f,mu1,mu2,METHOD);
   z = reshape(z,n^2,1);
   invM(:,i)=z;
end

end