/* This class implements the 1D Poisson Matrix operators, i.e., 
 * matrix-vector multiply and psolve.
 */

public class Poisson1DBlockJacobiPrecond implements Preconditioner {
    
    public Poisson1DBlockJacobiPrecond (int n) 
    {
	this.n = n;
	final int high = blockSize-1;
	factoredBlock = new double [0:high,0:high];
	factoredBlockT = new double [0:high,0:high];
	factoredBlock.set(0); // unecessary
	factoredBlock[0,0] = 2;
	factoredBlock[0,1] = -1;
	foreach (i in [1:high-1]) {
	    factoredBlock[i[1],i[1]] = 2;
	    factoredBlock[i[1],i[1]-1] = -1;
	    factoredBlock[i[1],i[1]+1] = -1;
	}
	factoredBlock[high,high-1] = -1;
	factoredBlock[high,high] = 2;
	Matrix.cholesky(factoredBlock);
	factoredBlockT.copy(factoredBlock.permute([2,1]));
    } 

    /* Assumes minvx and x are aligned in coordinate space */
    public void psolve(double [1d] minvx, double [1d] x)
    {
      minvx.copy(x);
      minvx = minvx.translate([0]-minvx.domain().min());
      for (int i = 0; i < n; i += blockSize) {
	Matrix.triangSolveLower(factoredBlockT,
			     minvx.restrict([i:i+blockSize-1]));
	Matrix.triangSolveUpper(factoredBlock,
			     minvx.restrict([i:i+blockSize-1]));
      }
    }

    /* Internal State */
    int n;
    static protected double [2d] factoredBlock;
    static protected double [2d] factoredBlockT;  // Simple, but inefficient
    static final protected int blockSize = 20;

    /* Test code */
    public static void tester (String [] args) 
    {
	if (args.length < 1) {
	    printUsage();
	    return;
	}
	int sz = Integer.parseInt(args[0]);

	// Allocate source and destination vectors
	double [1d] vec1 = new double [0:sz-1];
	double [1d] vec2 = new double [0:sz-1];
	vec1.set(1.0);

	// Test: Poisson matrix on psolve
	Poisson1DBlockJacobiPrecond p = new Poisson1DBlockJacobiPrecond(sz);
       	System.out.println("Poisson matrix: \n" + p);
	System.out.println("Source vector v1: \n" + 
			   Vector.toString(vec1) + "\n");
	p.psolve(vec2, vec1);	
	System.out.println("Result vector p.psolve(v1): " + 
			   Vector.toString(vec2));

	// Test: Poisson matrix on psolve using a subvector
	vec1 = new double [0:2*sz-1];
	vec1.set(1.0);
	vec2 = new double [0:2*sz-1];
	System.out.println("Source vector v1: \n" + 
			   Vector.toString(vec1) + "\n");
	vec1 = vec1.restrict(vec1.domain().shrink(sz/2));
	p.psolve(vec2, vec1);	
	System.out.println("Result vector p.psolve(v1): " + 
			   Vector.toString(vec2));
    }

    private static void printUsage() {
	System.out.println("Poisson1D <n> \n for a test run on an nxn matrix");
    }
}