package com.imsl.test.example.math;

import com.imsl.math.*;

/**
 * <p>
 * Solves the Poisson equation using the second Householder 
 * implementation and preconditioning. </p>
 *
 * The coefficient matrix in this example corresponds to the five-point
 * discretization of the 2-d Poisson equation with the Dirichlet boundary
 * condition. Assuming the natural ordering of the unknowns, and moving all
 * boundary terms to the right hand side, we obtain a block tridiagonal matrix.
 * (Consider the tridiagonal matrix \(T\) which has the value 4.0 down the main
 * diagonal and -1.0 along the upper and lower co-diagonals. Then the
 * coefficient matrix is the block tridiagonal matrix consisting of \(T\)'s down
 * the main diagonal and \(-I\), the identity matrix, along the upper and lower
 * co-diagonals.)
 * </p>
 * <p>
 * Discretizing on a 20 x 20 grid implies that the coefficient matrix is 400 x
 * 400. In the solution, the second Householder implementation is selected and
 * we choose to update the residual vector by direct evaluation.
 * </p>
 *
<p>
 * Preconditioning is used with the preconditioning matrix being a diagonal
 * matrix with 4.0 down the main diagonal and -1.0 along the upper and lower
 * co-diagonals. This <code>preconditioner</code> method solves this tridiagonal
 * matrix.
 * </p>
 *
 * @see <a href="GenMinResEx6.java">Code</a>
 * @see <a href="doc-files/GenMinResEx6.out">Output</a>
 */

public class GenMinResEx6 implements GenMinRes.Preconditioner {

    private double precondA[], precondB[], precondC[];

    // Creates a new instance of GenMinResEx6 
    public GenMinResEx6(int n) throws Exception {
        precondA = new double[n];
        precondB = new double[n];
        precondC = new double[n];
        // Define the preconditioning matrix                    
        for (int i = 0; i < n; i++) {
            precondA[i] = 4.0;
            precondB[i] = -1.0;
            precondC[i] = -1.0;
        }
    }

    /**
     * Obtains the multiplication of the matrix <code>a</code> and the input
     * <code>p</code>. The result is returned in <code>z</code>.
     *
     * @param p a <code>double</code> array with
     * <code>p.length</code>=<code>a[0].length</code>
     * @param z a <code>double</code> array
     */
    @Override
    public void amultp(double p[], double z[]) {
        int m = z.length;
        int n = (int) Math.sqrt(m);
        // Multiply by diagonal blocks           
        for (int i = 0; i < m; i++) {
            z[i] = 4.0 * p[i];
        }
        for (int i = 0; i < m - 2; i++) {
            z[i] -= p[i + 1];
        }
        for (int i = 0; i < m - 2; i++) {
            z[i + 1] -= p[i];
        }
        // Correct for terms not properly in block diagonal
        for (int i = n - 1; i < m - n; i = i + n) {
            z[i] += p[i + 1];
            z[i + 1] += p[i];
        }
        // Do the super and subdiagonal blocks, the -I's     
        for (int i = 0; i < m - n; i++) {
            z[i] -= p[i + n];
        }
        for (int i = 0; i < m - n; i++) {
            z[i + n] -= p[i];
        }
    }

    /**
     * Solve the tridiagonal preconditioning matrix problem for z.
     */
    public void preconditioner(double r[], double z[]) {
        int n = z.length;
        double w[] = new double[n];
        double v[] = new double[n];
        double u[] = new double[n];
        w[0] = precondA[0];
        v[0] = precondC[0] / w[0];
        u[0] = r[0] / w[0];
        for (int i = 1; i < n; i++) {
            w[i] = precondA[i] - precondB[i] * v[i - 1];
            v[i] = precondC[i] / w[i];
            u[i] = (r[i] - precondB[i] * u[i - 1]) / w[i];
        }
        z[n - 1] = u[n - 1];
        for (int j = n - 2; j >= 0; j--) {
            z[j] = u[j] - v[j] * z[j + 1];
        }
    }

    public static void main(String args[]) throws Exception {
        int n = 400;
        double b[] = new double[n];
        double xguess[] = new double[n];

        GenMinResEx6 atp = new GenMinResEx6(n);

        // Construct a GenMinRes object
        GenMinRes gnmnrs = new GenMinRes(n, atp);
        // Set right hand side and initial guess to ones           
        for (int i = 0; i < n; i++) {
            b[i] = 1.0;
            xguess[i] = 1.0;
        }
        gnmnrs.setGuess(xguess);
        gnmnrs.setMethod(gnmnrs.SECOND_HOUSEHOLDER);
        gnmnrs.setResidualUpdating(gnmnrs.DIRECT_AT_RESTART_ONLY);
        // Solve Ax = b
        gnmnrs.solve(b);
        int iterations = gnmnrs.getIterations();
        System.out.println("The number of iterations used = " + iterations);
        double resnorm = gnmnrs.getResidualNorm();
        System.out.println("The final residual norm is " + resnorm);
    }
}