#include <stdio.h>
#include <math.h>
#include <vector>
#include <iostream>
using namespace std;

/* takes one step of jacobi's method on 1d Poisson
   n+1 is the number of total grid points, including boundary, which
   are assumed to be 0 
   overwrites input vector x with the new result.
   b is the right-hand side ( b_i = h * f_i ) */
int jacobi_step(const vector<double> &b,
		vector<double> &x)
{
  vector<double> xold(x);

  int i;

  for (i=0;i<b.size();i++) {
    xold[i] = x[i];
  }

  /* write the next iteration into x */
  for (i=1;i<b.size()-1;i++) {
    x[i] = 0.5 * ( b[i] + xold[i-1] + xold[i+1] );
  }

  return 0;
}

void residual( const vector<double> &b, 
	       const vector<double> &x,
	       vector<double> &r )

{
  int n = b.size() - 1;
  int i;
  r[0] = 0.0;
  for (i=1;i<n;i++) {
    r[i] = 2.0 * x[i] - x[i-1] - x[i+1] - b[i];
  }
  r[n] = 0.0;

  return;
}

void dumb_restriction( const vector<double> &r ,
		       vector<double> &Rr )
{
  int n = (r.size()-1)/2;
  int i;

  Rr[0] = 0.0;
  for (i=1;i<n;i++) {
    Rr[i] = r[2*i];
  }
  Rr[n] = 0.0;

  return;
}

/* restricts from level i to level i - 1 */
void restriction( const vector<double> &r ,
		  vector<double> &Rr)
{
  int n = (r.size()-1)/2;
  int i;

  Rr[0] = 0.0;
  for (i=1;i<n;i++) {
    Rr[i] = 2*r[2*i] + r[2*i+1] * r[2*i-1];
  }
  Rr[n] = 0.0;

  return;
}

void injection( const vector<double> &a ,
		vector<double> &Aa)
{
  int n = a.size()-1;
  int i;
  
  Aa[0] = 0.0;
  for (i=1;i<=n;i++) {
    Aa[2*i-1] = 0.5 * (a[i] + a[i-1]);
    Aa[2*i] = a[i];
  }

  return;
}

int mgv(const int level,   /* n = 2^(level) + 1 */
	const vector<double> &b,
	vector<double> &x)
{
  int n1 = b.size();
  int n2 = (n1-1)/2 + 1;
  vector<double> r(n1);
  vector<double> R_r(n2);
  vector<double> d(n1);
  vector<double> a(n2);
  
  FILE *f;
  char filename[80];

  if (level == 1) {
    x[1] = b[1] / 2.0;
  }
  else {
    int i;
    int num_jacobi = 15;
    
    /* presmoothing by some jacobi steps */
    for (i=0;i<num_jacobi;i++) {
      jacobi_step(b,x);
    }
 
    /* compute residual by r = Ax - b */
    residual(b,x,r);

    /* apply 4R to r */
    restriction(r,R_r);

    /* a = mgv(level-1,4R(r),0) */
    for (i=0;i<a.size();i++) {
      a[i] = 0.0;
    }
    mgv(level-1,R_r,a);

    /* d = In(a) */
    injection(a,d);

    /* x = x - d */
    for (i=1;i<x.size()-1;i++) {
      x[i] -= d[i];
    }

    /* postsmoothing */
    for (i=0;i<num_jacobi;i++) {
      jacobi_step(b,x);
    }
  }
}


/* iterates jacobi's method until convergence or maximum 
   iteration count is exceeded.  Returns number of steps taken,
   with failure indicated by the return value being = max_iter */
int jacobi(const vector<double> &b,
	   const double tol,
	   const int max_iter,
	   vector<double> &x)
{
  double resid;
  double tmp;
  int i,j;
  char filename[80];
  FILE *f;
  int n = b.size()-1;
  double h = 1.0 / n;
  i = 0;

  resid = 2 * tol;

  while ( (i<max_iter) && (resid > tol) ) {
    vector<double> xprev(x);

    jacobi_step(b,x);


    resid = fabs( xprev[1] - x[1] ) / fabs( x[1] );
    for (j=2;j<b.size();j++) {
      if ( fabs( xprev[j] - x[j] ) / fabs( x[j] ) > resid ) {
	resid = fabs( xprev[j] - x[j] ) / fabs( x[j] );
      }
    }

    i++;

//     if (i % 2 == 0) {
//       sprintf(filename,"jacstep%d.dat",i);
//       f = fopen(filename,"w");
//       for (j=0;j<=n;j++) {
// 	fprintf(f,"%e\t%e\n",j*h,x[j]);
//       }
//       fclose(f);
//     }

  }

  return i;

}

int mgvcycles(const int level,
	      const vector<double> &b,
	      const double tol,
	      const int max_iter,
	      vector<double> &x)
{
  double resid;
  double tmp;
  int i,j;
  char filename[80];
  FILE *f;
  int n = b.size()-1;
  double h = 1.0 / n;
  i = 0;


  resid = 2 * tol;

  while ( (i<max_iter) && (resid > tol) ) {
    vector<double> xprev(x);

    mgv(level,b,x);


    resid = fabs( xprev[1] - x[1] ) / fabs( x[1] );
    for (j=2;j<b.size();j++) {
      if ( fabs( xprev[j] - x[j] ) / fabs( x[j] ) > resid ) {
	resid = fabs( xprev[j] - x[j] ) / fabs( x[j] );
      }
    }

    i++;

//     if (i % 2 == 0) {
 //       sprintf(filename,"mgvstep%d.dat",i);
//       f = fopen(filename,"w");
//       for (j=0;j<=n;j++) {
// 	fprintf(f,"%e\t%e\n",j*h,x[j]);
//       }
//       fclose(f);
//     }

  }

  return i;

}

/* uses v-cycles on one level to get initial guess on next finer level */
int fmg(const int max_level,
	const vector<double> &b,
	vector <double> &x)
{
  int n;
  int num_v_cycles = 30;
  int i,j;

  vector<vector<double> > bs(max_level);
  vector<vector<double> > xs(max_level);

  bs[max_level-1] = b;
  xs[max_level-1] = x;

  n = x.size()-1;

  for (i=max_level-1;i>=1;i--) {
    n/=2;
    bs[i-1] = vector<double>(n+1);
    xs[i-1] = vector<double>(n+1);
  }

  for (i=max_level-1;i>=1;i--) {
    restriction(bs[i],bs[i-1]);
    restriction(xs[i],xs[i-1]);
  }

  for (i=0;i<max_level-1;i++) {
    n = bs[i].size();
    mgvcycles(i+1,bs[i],1.e-5,n*n,xs[i]);
//     for (j=0;j<num_v_cycles;j++) {
//       mgv(i+1,bs[i],xs[i]);
//     }

    injection(xs[i],xs[i+1]);

  }

  for (j=0;j<num_v_cycles;j++) {
    mgv(i+1,b,xs[max_level-1]);
  }

  x = xs[max_level-1];

}





int main() {
  int level = 5;
  int n = 1 << level;
  double tol = 1.e-4;

  vector<double> x(n+1);
  vector<double> b(n+1);
  vector<double> xtrue(n+1);

  int i,j;
  double h;     /* spacing of one interval */
  double pt;
  int iter;
  double error;

  h = 1.0 / n;
  
  /* set up right hand side */
  for (i=0;i<=n;i++) {
    pt = i * h;
    b[i] = h * h  * 16 * M_PI * M_PI * sin(4*M_PI * pt);
    xtrue[i] = sin(4*M_PI*pt);
//     if (pt < 0.5) {
//       b[i] = h * h * pt * 16 * sin(4*M_PI * pt);
//     }
//     else {
//       b[i] = h * h * 16 * (1.0-pt);
//     }

  }

  

  for (i=0;i<=n;i++) {
    x[i] = 0.0;
  }  
  
  FILE *f;
  
  iter = jacobi(b,tol,n*n*n,x);      

  f = fopen("soljac.dat","w");
  for (i=0;i<=n;i++) {
    pt = i * h;
    fprintf(f,"%f\t%f\n",pt,x[i]);
  }
  fclose( f );

  error = 0.0;
  for (i=1;i<n;i++) {
    error += h * pow( xtrue[i] - x[i] , 2 );
  }
  error = sqrt( error);
  printf( "error=%e\n",error );

//   for (i=0;i<=n;i++) {
//     x[i] = 0.0;
//   }  

//   iter = mgvcycles(level,b,tol,n*n*n,x);

//   f = fopen("solmgv.dat","w");
//   for (i=0;i<=n;i++) {
//     pt = i * h;
//     fprintf(f,"%f\t%f\n",pt,x[i]);
//   }
//   fclose( f );
  
//   error = 0.0;
//   for (i=1;i<n;i++) {
//     error += h * pow( xtrue[i] - x[i] , 2 );
//   }
//   error = sqrt( error);
//   printf( "error=%e\n",error );


  for (i=0;i<=n;i++) {
    x[i] = 0.0;
  }  

  fmg(level,b,x);

  error = 0.0;
  for (i=1;i<n;i++) {
    error += h * pow( xtrue[i] - x[i] , 2 );
  }
  error = sqrt( error);
  printf( "error=%e\n",error );

  f = fopen("solfmg.dat","w");
  for (i=0;i<=n;i++) {
    pt = i * h;
    fprintf(f,"%f\t%f\n",pt,x[i]);
  }
  fclose( f );
  
  f = fopen("soltrue.dat","w");
  for (i=0;i<=n;i++) {
    pt = i * h;
    fprintf(f,"%f\t%e\n",pt,xtrue[i]);
  }
  
  return 0;
}