// Jacobi.java
 

package dogma.examples.mpij;
 

import mpij.*;
import java.util.*;
import java.io.*;
 

public class Jacobi extends MPIApplication {

    public void MPIMain(String[] args) {
        int size;
        int MAXX;
        int MAXY;
        int MAXSTEPS;

        int iproc, nproc;
        int unbor, dnbor;
        int i, j;
        int nrows, ncols;
        int steps;

        Status status;
        double t0;
        Random rand;

        double[] time;
        double[][] iplate;
        double[][] oplate;
        double[][] tmp;

        if (args.length != 1) {
            out.println("Usage dogma.examples.mpij.Jacobi <size>");
            return;
        }
        size = Integer.parseInt(args[0]);
        MAXX = size;
        MAXY = size;
        MAXSTEPS = 128;
        status = new Status();
        rand = new Random();
        time = new double[1];

        try {
            MPI.init();
            iproc = MPI.COMM_WORLD.rank();
            nproc = MPI.COMM_WORLD.size();
            if (size % nproc != 0) {
                out.println("size must be divisible by the "
                    +"number of processes");
                MPI.finalize();
            }

            ncols = MAXY;
            nrows = MAXX / nproc;

            out.println("nrows "+ nrows +" ncols "+ ncols);

            iplate = new double[nrows + 2][ncols];
            oplate = new double[nrows + 2][ncols];

            //Determine my up and down neighbors
            unbor = (iproc + nproc - 1) % nproc;
            dnbor = (iproc + 1) % nproc;
            if (iproc == 0)
                unbor = -1;
            if (iproc == nproc - 1)
                dnbor = -1;

            // Initialize the array using random numbers
            for (i = 0; i < nrows + 2; i++) {
                for (j = 0; j < MAXY; j++) {
                    iplate[i][j] = rand.nextFloat();
                }
            }
            MPI.COMM_WORLD.barrier();

            // Now proceed with the Jacobi algorithm
            t0 = MPI.wtime();

            for (steps = 0; steps < MAXSTEPS; steps++) {

                // Exchange boundary conditions with my neighbors
                if (unbor != -1) {
                    MPI.COMM_WORLD.send(iplate[1], 0, ncols, MPI.DOUBLE,
                        unbor, 1);
                    MPI.COMM_WORLD.recv(iplate[0], 0, ncols, MPI.DOUBLE,
                        unbor, 1, status);
                }

                if (dnbor != -1) {
                    MPI.COMM_WORLD.send(iplate[nrows-2], 0, ncols, MPI.DOUBLE,
                        dnbor, 1);
                    MPI.COMM_WORLD.recv(iplate[nrows-1], 0, ncols, MPI.DOUBLE,
                        dnbor, 1, status);
                }

                // Now compute the 5 point stencil
                for (i = 1; i < nrows+1; i++) {
                    for (j = 1; j < ncols-1; j++) {
                        oplate[i][j] = (double)(4 * iplate[i][j] +
                                iplate[i-1][j] +
                                iplate[i+1][j] +
                                iplate[i][j-1] +
                                iplate[i][j+1])/8.0f;
                    }
                }

                // Now copy the new results into the current plate array
                tmp = oplate;
                oplate = iplate;
                iplate = tmp;
            }

            time[0] = MPI.wtime() - t0;
            out.println("My Elapsed Time is " + time[0] + " sec");

            // determine the maximum elapsed time
            MPI.COMM_WORLD.reduce(time, 0, time, 0, 1, MPI.DOUBLE, MPI.MAX, 0);
            if (iproc == 0) {
                out.println("Max Total Time is " + time[0] + " sec");
            }

            MPI.finalize();
        } catch (MPIException me) {
            me.printStackTrace();
        }
    }
}