#ifndef __CG_H
#define __CG_H

// Iterative template routine -- conjugate gradient method
//
// CG solves the linear system Ax=b using conjugate gradients.
// The returned value indicates convergence within
// max_iter iterations (return value 0)
// or no convergence within max_iter iterations (return value 1)
// Upon successful return (0), the output arguments have the
// following values:
//        x: computed solution
// mat_iter: number of iterations to satisfy the stopping criterion
//      tol: residual after the final iteration

template <class MATRIX, class VECTOR, class REAL>
int
CG (const MATRIX & A, VECTOR & x, const VECTOR & b,
	    int & max_iter, REAL & tol)
{
  REAL resid;
  VECTOR p(b.size ());
  VECTOR q(b.size ());
  REAL alpha, beta, rho, rho_1=0;
  REAL normb = b.l2norm();
  VECTOR r = b - A * x;

  if (normb == 0.0) normb = 1;
  resid = r.l2norm() / normb;
    
  // IF YOU WANT OUTPUT
  cout << "CG " << 0 << ": " << resid; cout.flush ();
  
  if (resid <= tol)
    {
      tol = resid;
      max_iter = 0;
      return 0;
    }
  
  for (int i=1; i<=max_iter; ++i)
    {
      rho = inner_product (r, r);
      if (i == 1)
	{
	  p = r;
	}
      else
	{
	  beta = rho / rho_1;
	  p = r + beta * p;
	}

      q = A * p;

      alpha = rho / inner_product (q, p);

      x += alpha * p;
      r -= alpha * q;

      resid = r.l2norm() / normb;

      // IF YOU WANT OUTPUT
      cout << "\r" << "CG " << i << ": " << resid << "       "; cout.flush();

      if (resid <= tol)
	{
	  tol = resid;
	  max_iter = i;
          cout << endl;
	  return 0;
	}
      
      rho_1 = rho;
    }
  
  tol = resid;
  cout << endl;
  return 1;
}

#endif