// two-demo.C : demonstration program for using twodimensional grafix tools
// to drive the routines sor and mglin
//  wolf 11/96

#include "toolbox.h"
#include <math.h>

#define NRANSI
extern "C" {
#include "nr.h"
#include "nrutil.h"
}

// displayable arrays must be defined globally !!
double **a,**b,**c,**d,**e; // coefficients for sor
double **f; // source term
double **err; // error
double **u; // solution 

float fwidth = 1; // must be initialized 

void init_arrays(int n) {
  for (int i=1; i <= n; i++)
    for (int j=1; j <= n; j++) {
      a[i][j] = b[i][j] = c[i][j] = d[i][j] = 1.0;
      e[i][j] = -4.0;
      f[i][j] = u[i][j] = 0.0;
    }
}

void source(int n) { // compute source term as function of fwidth
  int midl = n/2 + 1, i = midl, j = midl;
  for (int k = 0; k < fwidth; k++) { 
    f[i][j] = -2000.0/((n-1)*(n-1));  // use a convinient norming 
    u[i][j] = -2000; // should be 0 for sor
    j++; if (j == n) { j = 2; i++; }
  }
}

void error(int n) {
  // compute the approx. error array, f is the source term 
  // *** this is the error of the difference eq., not the PDE !! ***
  for (int i=1; i <= n; i++) for (int j=1; j<= n; j++) err[i][j] = f[i][j];
  for (int i=2; i < n; i++) for (int j=2; j < n; j++) 
    err[i][j] -= u[i+1][j] + u[i-1][j] + u[i][j+1] + u[i][j-1] - 4.0*u[i][j];
}

// two solvers for the elliptic Poisson equation are installed :
// 1. Successive Overrelaxation (SOR) uses NR function "sor"
// 2. Full Multigrid Algorithm (FMG) uses NR function "mglin"

void sor_wrapper(int nx, int ny) { // wrapper for sor
  if (nx != ny) { printf("wrong dimension nx != ny"); exit(1); }
  int n = nx;
  init_arrays(n); 
  source(n);  
  double rjac = cos(M_PI/n); 
  printf("calling sor for n = %d %f %f\n",n,fwidth,rjac);
  sor(a,b,c,d,e,f,u,n,rjac);
  error(n);
}

void mglin_wrapper(int nx, int ny) { // wrapper for sor
  if (nx != ny) { printf("wrong dimension nx != ny"); exit(1); }
  int n = nx;
  for (int i=1; i<= n; i++)
    for (int j=1; j<= n; j++) f[i][j] = u[i][j] = 0.0;
  source(n);
  // printf("calling mglin for n = %d\n",n);
  mglin(u,n,2);
  error(n);
}

int main(void) {
  int n = 33;
  a = dmatrix(1, n, 1, n);
  b = dmatrix(1, n, 1, n);
  c = dmatrix(1, n, 1, n);
  d = dmatrix(1, n, 1, n);
  e = dmatrix(1, n, 1, n);
  f = dmatrix(1, n, 1, n);
  u = dmatrix(1, n, 1, n);   
  err = dmatrix(1, n, 1, n);  

  // declare which arrays should be displayable, together with description
  // and initial z-magnification factor (empty means auto detect), 
  // !!! do not forget the closing '0' at the end of the lists !!!
  disp_array *da[] = { new disp_array(u,"u array auto"), // auto factor
		       new disp_array(u,"u array zinit",zinit),
		       new disp_array(u,"u array fixed",0.1), 
		       new disp_array(f,"source"), // auto factor
		       new disp_array(err,"error"),
		       0 }; // !!!

  // the list of solvers for the equation
  two_solver *sv[] = { new two_solver(mglin_wrapper,"mglin"),
		       new two_solver(sor_wrapper,"sor"), 
		       0 }; // !!!

  parameter *scr[] = { 
    new parameter("source width", &fwidth, 0, 200), 
    0 }; // !!!!

  // start the display loop
  two_show("two-demo",n,n,da,sv, scr); // scr is optional (may be omitted) 

  free_dmatrix(u,1,n,1,n);
  free_dmatrix(f,1,n,1,n);
  free_dmatrix(e,1,n,1,n);
  free_dmatrix(d,1,n,1,n);
  free_dmatrix(c,1,n,1,n);
  free_dmatrix(b,1,n,1,n);
  free_dmatrix(a,1,n,1,n);

  return 0;
}