#include <iostream>
#include <iomanip>

typedef void (*MatrixAction) (double*, double*, int n);

// y = alpha*x (x,y - vectors, alpha - scalar)
inline void vector_copy(double *y, double *x, int n, double alpha = 1.0)
{
   for(int i=0; i<n; i++)
      y[i] = alpha*x[i];
}

// y = y + alpha*x (x,y - vectors, alpha - scalar)
inline void vector_add(double *y, double *x, int n, double alpha = 1.0)
{
   for(int i=0; i<n; i++)
      y[i] += alpha*x[i];
}

inline double vector_inner_product(double *x, double *y, int n)
{
   double result=0;
   for(int i=0; i<n; i++)
      result += x[i]*y[i];
   return result;
}
   

// A procedure implementing the Conjugate Gradient algorithm for
// solving the equation Ax=b, where the parameters are:
// A   - a function implementic the action of the symmetric and positive
//       definte matrix (called A again in our equation),
// x   - a vector containing the RHS on entering and the solution on exiting
// n   - the dimention of the vector x (and the sqare matrix A)
// eps      = the tolerance
// max_iter = the maximum number of iterations allowed
int solveCG(MatrixAction A, double *x, int n,
            double eps = 0.1e-16, int max_iter = 20000)
{
   //We can do without some of these variables, if we reuse others
   //  I have not tried to optimize the code, but to write it as
   //  clear as possible.
   
   double *r = new double[n];
   double *p = new double[n];
   double *temp = new double[n];
   double alpha, beta, r_norm;

   int iter = 0;

   //Initialize (we assume x_0=0 is our initial guess):
   vector_copy(r,x,n);    // r_0 = b (the right hand side)
   vector_copy(x,x,n,0);  // x_0 = 0  (and so, r_0 = b = b - Ax_0, which is what we want
   vector_copy(p,r,n);    // p_0 = r_0
   r_norm = vector_inner_product(r,r,n);

   while ((r_norm > eps) && (iter < max_iter)) {
      iter++;
      
      //Step 1:
      A(temp,p,n);  //temp = Ap(matrix vector product)
      alpha = r_norm / vector_inner_product(temp,p,n);  // alpha = (r_n,r_n) / (Ap_n,p,n)
      vector_add(x,p,n,alpha); // x_(n+1) = x_n + alpha p_n
      
      //Step 2:
      vector_add(r,temp,n,-alpha); //r_(n+1) = r_n - alpha A p_n

      //Step 3:
      beta = r_norm; //(r_n,r_n)
      r_norm = vector_inner_product(r,r,n); //(r_(n+1),r_(n+1))
      beta = r_norm / beta; //  (r_(n+1),r_(n+1)) / (r_n,r_n)
      
      vector_copy(temp,r,n);
      vector_add(temp,p,n,beta); //p_(n+1) = (temp) = r_(n+1) + beta p_n
      vector_copy(p,temp,n);
   }

   delete p;
   delete r;
   delete temp;

   return iter;
}