public class Multigrid {
    public Multigrid nextCoarser;
    public RectDomain<2> D, interiorD;
    public Domain<2> borderD; // D - interiorD
    public Domain<2> red, black; // red and black checkboard of interiorD
    public double[2d] phi, res, rhs;
    public double h, k;
    public int x, y;

    public Multigrid(int x0, int y0, double h0) {
	x = x0;
	y = x0;
	h = h0;
	k = 6.0 * h * h; // used in stencils
	System.out.println(Ti.thisProc() + "Multigrid construction: x = " + x + ", y = " + y +
			   ", h = " + h);
	D = [0 : x + 2, 0 : y + 2];
	interiorD = [1 : x + 1, 1 : y + 1];
	borderD = D - interiorD;
	{
	    RectDomain<2> red1 = [[1, 1] : [x + 1, y + 1] : [2, 2]];
	    RectDomain<2> red2 = [[2, 2] : [x + 1, y + 1] : [2, 2]];
	    RectDomain<2> black1 = [[1, 2] : [x + 1, y + 1] : [2, 2]];
	    RectDomain<2> black2 = [[2, 1] : [x + 1, y + 1] : [2, 2]];

	    red = red1 + red2;
	    black = black1 + black2;
	}
	phi = new double[D];
	res = new double[D];
	rhs = new double[D];
	if (x > 3 || y > 3) nextCoarser = new Multigrid(x / 2, y / 2, h * 2.0);
    }

    public void relax() {
	if (Ti.thisProc() == 0) System.out.println("Relaxing");
	print(rhs);
	gsrb();
	gsrb();
	if (nextCoarser != null) {
	    foreach (p in nextCoarser.D)
		nextCoarser.phi[p] = 0.;
	    residual();
	    average();
	    nextCoarser.relax();
	    interpolate();
	}
	gsrb();
	gsrb();
	print(phi);
	if (Ti.thisProc() == 0) System.out.println("Done");
    }

    public void gsrb() {
	boundary(phi);
        gsrb_color(red);        
        gsrb_color(black);        
    }
    public inline void gsrb_color(Domain<2> D) {
	    foreach (q in D)
		res[q] = ( (phi[n(q)] + phi[s(q)] + phi[e(q)] + phi[w(q)]) * 4.0
			   + (phi[ne(q)] + phi[se(q)] + phi[nw(q)] + phi[sw(q)])
			   - 20.0 * phi[q] - k * rhs[q] ) * 0.05;
	    foreach (q in D)
		phi[q] += res[q];
    }

    public void boundary(double[2d] a) {
	foreach (p in borderD) a[p] = 0;
    }

    public void residual() {
	boundary(phi);
	foreach (q in interiorD) 
	    res[q] = ( (phi[n(q)] + phi[s(q)] + phi[e(q)] + phi[w(q)]) * 4.0
		       + (phi[ne(q)] + phi[se(q)] + phi[nw(q)] + phi[sw(q)])
		       - 20.0 * phi[q] - k * rhs[q]) / -k;
	boundary(res);
    }


    public void average() {
	foreach (q in nextCoarser.interiorD) {
	    Point<2> qfine = q + q;
	    nextCoarser.rhs[q] =
		.0625 * (4.0 * res[qfine] +
			 (res[e(qfine)] + res[w(qfine)] +
			  res[n(qfine)] + res[s(qfine)]) * 2.0 +
		res[ne(qfine)] + res[nw(qfine)] +
		res[se(qfine)] + res[sw(qfine)]);
	    }
    }

    public void interpolate() {
	double[2d] coarse = nextCoarser.phi;

	foreach (q in nextCoarser.interiorD) {
	    Point <2> q2 = q + q;
	    phi[sw(q2)] += .25 * (coarse[q] + coarse[sw(q)] +
			          coarse[w(q)] + coarse[s(q)]);
	    phi[s(q2)] += .5 * (coarse[q] + coarse[s(q)]);
	    phi[w(q2)] += .5 * (coarse[q] + coarse[w(q)]);
	    phi[q2] += coarse[q];
	}
	boundary(phi);
    }

    // Since the stencil is symmetric, don't go complaining if you
    // think north and south are backwards here.  It doesn't matter.
    // (Let's hope these are inlined.)
    public static Point<2> n(Point<2> p) { return p + [0, 1]; }
    public static Point<2> s(Point<2> p) { return p + [0, -1]; }
    public static Point<2> e(Point<2> p) { return p + [1, 0]; }
    public static Point<2> w(Point<2> p) { return p + [-1, 0]; }
    public static Point<2> ne(Point<2> p) { return p + [1, 1]; }
    public static Point<2> nw(Point<2> p) { return p + [-1, 1]; }
    public static Point<2> se(Point<2> p) { return p + [1, -1]; }
    public static Point<2> sw(Point<2> p) { return p + [-1, -1]; }

    public void print(double[2d] a) {
	if (Ti.thisProc() == 0) {
	    System.out.println("");
	    for (int j = y; j >= 0; j--) {
		// System.out.print("j = " + j + ":\t");
		for (int i = 0; i <= x; i++) System.out.print(a[[i, j]] + "\t");
		System.out.println("");
	    }
	}
    }

    public static void main(String[] args) {
	int x = 8, y = 8;

	Multigrid m = new Multigrid(x, y, 0.0625);

	for (int i = 1; i < x; i++)
	    for (int j = 1; j < y; j++)
		m.rhs[[i, j]] = i * i + j * j;
	// m.print(m.rhs);
	// m.print(m.phi);
	// m.gsrb();
	m.relax();
	// m.print(m.phi);
    }
  
}