/* This program demonstrates the Distributed Arrays (DA) concept and PETSc solvers. 
 Poisson equation of the form
   laplace (q) = f(x,y)
   0 <= x <= 1, 0 <= y <= 1
 Discretization is performed via central second-order finite difference scheme.
 Resulting linear system
  [A]{q} = {rhs}
 has a penta-diagonal symmetric which is solved by various solvers and preconditioners. 

Run the code with: 
./ex1 -Nx <nx> -Ny <ny> -bc <bc_type> -solver_type <solver_type>
Nx: 
	number of grid in x-direction
Ny: 
	number of grid in x-direction
bc: 
	1: Neumann
	2: Dirichlet
solver_type: 
	1: LU direct solver (requires MUMPS) 
	2: KSP GMRES solver 
	3: KSP GMRES solver with BoomerAMG preconditioner (requires Hypre package) 
	4: KSP GMRES solver with PCGAMG preconditioner (requires Hypre package) 


 By: Mohamad M. Nasr-Azadani
 Email: mmnasr@engineering.ucsb.edu
*/
static char help[] = "*********************************************************************************************************************************\n*********************************************************************************************************************************\nThis codes solves 2-D Poisson equation using DMDA arrays and PETSc solvers.\n\nBoundary conditions are either NEUMANN or DIRICHLET.\nTo run the code, simply pass the number of grid points in each direction and the boundary conditions (1:NEUMANN, 2:DIRICHLET),\ne.g. ./ex1 -Nx 10 -Ny 20 -bc 1 -solver_type 3\nOutput shows the norm1 and norm2 errors (E1 and E2) compared to analytical solution. The solver_type can be: 1: LU direct solver (needs MUMPS package)  2:KSP GMRES iterative solver 3: KSP GMRES solver with BoomerAMG preconditioner(needs Hypre package) 4: PETSc PCGAMG preconditioner.";

#include "petsc.h"
#include "petsctime.h"
#include "ex1.h"

int main(int argc,char **argv){

    PetscErrorCode ierr;
    /* Always start with this command to initialize your program */
    ierr = PetscInitialize(&argc,&argv,(char*)0,help);
    /* Always use this function to check error codes returned from PETSc function calls.
 It can be very useful when debugging your code */
    CHKERRQ(ierr);

    /* Check if enough input arguments are passed to the code from the command prompt. */
    if (argc < 9) {

        PetscPrintf(PCW, "%s\n", help);
        PetscPrintf(PCW, "Error! Not enough input parameters.\nPlease enter the number of grid points and boundary condition type (1:NEUMANN, 2:DIRICHLET): -Nx [] -Ny [] -bc [] -solver_type [], e.g.  \n");
        PetscPrintf(PCW, "  ./ex1 -Nx 10 -Ny 20 -bc 2 -solver_type 3\n");
        PetscPrintf(PCW, "*********************************************************************************************************************************\n");
        ierr = PetscFinalize();
        return 0;
    } /* if */

    /* Let's now get the number of processors (size) and current processor's rank */
    int size=-1;
    int rank=-1;
    ierr = MPI_Comm_size(PETSC_COMM_WORLD, &size);
    ierr = MPI_Comm_rank(PETSC_COMM_WORLD, &rank);

    /* 	Print number of processors being used (only on processor zero!). */
    PetscPrintf(PETSC_COMM_WORLD, "\nThis code runs on %d processor(s).\n", size);

    PetscInt Nx, Ny, bc;
    ierr = PetscOptionsGetInt(PETSC_NULL,"-Nx", &Nx, PETSC_NULL);CHKERRQ(ierr);
    ierr = PetscOptionsGetInt(PETSC_NULL,"-Ny", &Ny, PETSC_NULL);CHKERRQ(ierr);
    ierr = PetscOptionsGetInt(PETSC_NULL,"-bc", &bc, PETSC_NULL);CHKERRQ(ierr);
    PetscPrintf(PCW, "User provided input parameters: (Nx,Ny)=(%d,%d)\n", Nx, Ny);
    if (bc == NEUMANN) {
        PetscPrintf(PCW, "With Neumann boundary conditions.\n");
    }else {
        PetscPrintf(PCW, "With Dirichlet boundary condition.\n");
        bc = DIRICHLET;
    }
	//getchar(); 
    /****************************************************/
    /****************************************************/
    double dx = 1.0/(double)Nx;
    double dy = 1.0/(double)Ny;
    PetscInt n_ghosts = 1;
    PetscInt stencil_width = 1;
    DM DA2d;
    /* Let's create the parallel/serial DA layout which defines essentially the dimensions and layout of 2D array */
    /* Kind of like, double data[ny][nx] */
    /* Tip: If you are not sure of a certain parameter, use PETSC_DECIDE. Hopefully, it works! */
    ierr = DMDACreate2d(PETSC_COMM_WORLD, DMDA_BOUNDARY_NONE, DMDA_BOUNDARY_NONE, DMDA_STENCIL_STAR, Nx, Ny, PETSC_DECIDE, PETSC_DECIDE, stencil_width, n_ghosts, PETSC_NULL, PETSC_NULL, &DA2d); CHKERRQ(ierr);

    /* Let's get some information about the created DA layout */
    ierr = DMView(DA2d, 0);
    /****************************************************/
    /****************************************************/

    /* Associated with every DA, one can create two vectors, global and local */
    Vec RHS_vec_g;
    ierr = DMCreateGlobalVector(DA2d, &RHS_vec_g); CHKERRQ(ierr);

    /* Let's also create the solution and analytical solution vectors */
    Vec sol_vec, sol_an;
    ierr = VecDuplicate(RHS_vec_g, &sol_vec); CHKERRQ(ierr);
    ierr = VecDuplicate(RHS_vec_g, &sol_an); CHKERRQ(ierr);

    /* Now, setup the RHS associated with the Poisson equation */
    ierr = setup_RHS_vector(RHS_vec_g, DA2d, bc, dx, dy, Nx, Ny); CHKERRQ(ierr);
    //VecView(RHS_vec_g, 0);


    /****************************************************/
    /****************************************************/
    /* Now, let's create associated LHS matrix used with the current DA2d layout. */
    /* Note that, for each row in the matrix we will have (2*stencil_width + 1) = 5 nonzeros. */
    /* This can be thought as appropriate matrix structure to be used with second-order cental finite-difference schemes. */
    Mat LHS_mat;

    /* Not sure about matrix type, use MATMPIAIJ */
    ierr = DMCreateMatrix(DA2d,  MATMPIAIJ, &LHS_mat); CHKERRQ(ierr);

    /****************************************************/
    /****************************************************/

    ierr = setup_LHS_matrix(LHS_mat, DA2d, bc, dx, dy, Nx, Ny) ; CHKERRQ(ierr);

    /* Let's just see the matrix structure */
    //MatView(LHS_mat, 0);

    /****************************************************/
    /****************************************************/
    KSP solver;
    ierr = KSPCreate(PETSC_COMM_WORLD, &solver);
    ierr = KSPSetType(solver, KSPFGMRES); 
	
	PetscInt solver_type; 
    ierr = PetscOptionsGetInt(PETSC_NULL,"-solver_type", &solver_type, PETSC_NULL);CHKERRQ(ierr);
	if (solver_type == _LU_) {

	    PC pc;
	    ierr = KSPSetType(solver, KSPPREONLY);CHKERRQ(ierr);
	    ierr = KSPGetPC(solver, &pc);CHKERRQ(ierr);
	    ierr = PCSetType(pc, PCLU); CHKERRQ(ierr);
	    ierr = PCFactorSetMatSolverPackage(pc, MATSOLVERMUMPS);CHKERRQ(ierr);

	} else if (solver_type == _HYPRE_) {

	    PC pc;
	    ierr = KSPGetPC(solver, &pc);CHKERRQ(ierr);
		ierr = PCSetType(pc, PCHYPRE);CHKERRQ(ierr); 
		ierr = PCHYPRESetType(pc,"boomeramg");CHKERRQ(ierr); 

		ierr = PetscOptionsSetValue("-pc_hypre_boomeramg_max_levels", "25"); CHKERRQ(ierr);
		ierr = PetscOptionsSetValue("-pc_hypre_boomeramg_strong_threshold", "0.0"); CHKERRQ(ierr);
		ierr = PetscOptionsSetValue("-pc_hypre_boomeramg_relax_type_all", "SOR/Jacobi"); CHKERRQ(ierr);
	} else if (solver_type == _PETSC_GAMG_) { 

		PC pc;
		ierr = KSPGetPC(solver, &pc);CHKERRQ(ierr);
		ierr = PCSetType(pc, PCGAMG);CHKERRQ(ierr); 
	}
	
	/* For a NEUMANN B.C., the matrix has a nullspace condition. Therefore, we need these extra commands to let PETSc know about it. */
    if (bc == NEUMANN) {
        MatNullSpace nullsp;
        ierr = MatNullSpaceCreate(PETSC_COMM_WORLD, PETSC_TRUE, 0, PETSC_NULL, &nullsp); CHKERRQ(ierr);
        ierr = KSPSetNullSpace(solver, nullsp); CHKERRQ(ierr);
        ierr = MatNullSpaceDestroy(&nullsp);CHKERRQ(ierr);
    } /* if */
    ierr = KSPSetFromOptions(solver); CHKERRQ(ierr);
    ierr = KSPSetOperators(solver, LHS_mat, LHS_mat, SAME_NONZERO_PATTERN); CHKERRQ(ierr);


    //KSPView(solver, 0);

	PetscLogDouble t1, t2; 
	ierr = PetscGetTime(&t1); 
    /* Solve the linear system! */
    ierr = KSPSolve(solver, RHS_vec_g, sol_vec); CHKERRQ(ierr);
	ierr = PetscGetTime(&t2); 

    /* Let's now check the convergence status */
    KSPConvergedReason conv_flag;
    ierr = KSPGetConvergedReason(solver, &conv_flag); CHKERRQ(ierr);

    /* Return how many iterations it took to converge */
    PetscInt iters;
    ierr = KSPGetIterationNumber(solver, &iters); CHKERRQ(ierr);
    if (conv_flag < 0) {

        PetscPrintf(PETSC_COMM_WORLD, "Error! KSP solver did not converge! Divergence reason: %d (check PETSc manual for detailed instructions).\n", conv_flag);
        PetscPrintf(PETSC_COMM_WORLD, "KSP solver diverged with %d iterations\n", iters);
    } else {

        PetscPrintf(PETSC_COMM_WORLD, "KSP solver converged in %d iterations\n", iters);
		PetscPrintf(PCW, "Solution time: %2.14lf\n", (t2-t1)); 
    }
    /****************************************************/
    /****************************************************/

    ierr = setup_analytical_vector(sol_an, DA2d, bc, dx, dy, Nx, Ny) ;
    ierr = compute_error(sol_vec, sol_an, dx, dy) ;

    /****************************************************/
    /****************************************************/
    DMDestroy(&DA2d);
    VecDestroy(&RHS_vec_g);
    VecDestroy(&sol_vec);
    MatDestroy(&LHS_mat);
    KSPDestroy(&solver);

    /* Always finish your main program with this command */
    ierr = PetscFinalize();
    return 0;
}
/************************************************************************************************************************************************************/
/************************************************************************************************************************************************************/
/* This function sets up the matrix associated with the Poisson equation. 
   Second order central-finite difference stencil is used to discretize the Poisson equation on a uniform orthogonal grid:

            (i,j+1)
                |
                |
                |
  (i-1,j)-----(i,j)-----(i+1,j)
                |
                |
                |
            (i,j-1)
Two types of boundary condition can be used: 
1- Neumann     dq/dx, dq/dy = 0   on the boundaries (x=0,1, y=0,1) 
2- Dirichlet   q=0                on the boundaries (x=0,1, y=0,1) 

Boundary condition is implemeted by the "ghost node" concept. 
*/
PetscErrorCode setup_LHS_matrix(Mat A, DM DA2d, short int boundary_condition, PetscScalar dx, PetscScalar dy, PetscInt Nx, PetscInt Ny) {

    MatStencil cols[5], row; /* 5-point stencil */
    PetscScalar values[5];

    /* Discretization coefficients associated with all the current and neighboring nodes */
    PetscScalar ac0 = 2.0/(dx*dx) + 2.0/(dy*dy); /* (i,j)    */
    PetscScalar al = -1.0/(dx*dx);               /* (i-1,j)  */
    PetscScalar ar = -1.0/(dx*dx);               /* (i+1,j)  */
    PetscScalar at = -1.0/(dy*dy);               /* (i,j+1)  */
    PetscScalar ab = -1.0/(dy*dy); 				 /* (i,j-1)  */

    /* Get the index range of the 2D array on current processor */
    PetscInt Is, Js;
    PetscInt nx, ny;
    PetscErrorCode ierr;
    ierr = DMDAGetCorners(DA2d, &Is, &Js, PETSC_NULL, &nx, &ny, PETSC_NULL); CHKERRQ(ierr);
    PetscInt Ie = Is + nx;
    PetscInt Je = Js + ny;

    PetscScalar sign;
    if (boundary_condition == NEUMANN) {

        sign = +1.0;
    } else if (boundary_condition == DIRICHLET) {

        sign = -1.0;
    } else {

        PetscPrintf(PCW, "Error. Unknown boundary condition type.\n");
    }

    /* Go over all the nodes within the current processor range, i.e.
  j = Js to Je-1 (y-direction)
   i = Is to Ie-1 (x-direction)
  and set the non-zero values associated with the discretization used.
 */
    int i, j;
    for (j=Js; j<Je; j++) {
        for (i=Is; i<Ie; i++) {

            /* Number of nonzero elements in the row */
            PetscInt nnz=0;

            /* reset the value of ac */
            PetscScalar ac = ac0;

            /* Neighbor to the left (i-1,j) */
            if (i > 0) {
                cols[nnz].i = i-1;
                cols[nnz].j = j;
                values[nnz] = al;
                nnz++;
            } else {

                ac += sign*al;
            }

            /* Neighbor to the right (i+1,j) */
            if (i < Nx-1) {
                cols[nnz].i = i+1;
                cols[nnz].j = j;
                values[nnz] = ar;
                nnz++;

            } else {

                ac += sign*ar;
            }

            /* Neighbor on the top (i,j+1) */
            if (j < Ny-1) {
                cols[nnz].i = i;
                cols[nnz].j = j+1;
                values[nnz] = at;
                nnz++;
            } else {

                ac += sign*at;
            }

            /* Neighbor at the bottom (i,j-1) */
            if (j > 0) {
                cols[nnz].i = i;
                cols[nnz].j = j-1;
                values[nnz] = ab;
                nnz++;
            } else {

                ac += sign*ab;
            }

            cols[nnz].i = i;
            cols[nnz].j = j;
            values[nnz] = ac;
            nnz++;

            /* Geometric index of current node, i.e. (i,j) */
            row.i = i;
            row.j = j;

            //PetscPrintf(PCW, "(i,j)=(%d,%d) nnz:%d\n", i, j, nnz);
            /* Insert the non-zero values for currrent row in the matrix */
            ierr = MatSetValuesStencil(A, 1, &row, nnz, cols, values, INSERT_VALUES); CHKERRQ(ierr);
        }
    }

    /* ALWAYS call these two functions to let petsc know to finalize the Matrix. */
    ierr = MatAssemblyBegin(A, MAT_FINAL_ASSEMBLY); CHKERRQ(ierr);
    ierr = MatAssemblyEnd(A, MAT_FINAL_ASSEMBLY); CHKERRQ(ierr);

    return 0;
}
/************************************************************************************************************************************************************/
/************************************************************************************************************************************************************/

/* This function sets up the RHS vector associated with the Poisson equation. 
Depending on the boudary condition, different RHS function is automatically chosen, i.e. 

1- Neumann BC, rhs   = (2pi)*(2pi) cos(2pix)*cos(2piy) 
2- Dirichlet BC, rhs = (2pi)*(2pi) sin(2pix)*sin(2piy) 
*/
PetscErrorCode setup_RHS_vector(Vec b, DM DA2d, short int boundary_condition, PetscScalar dx, PetscScalar dy, PetscInt Nx, PetscInt Ny) {

    /* Get the index range of the 2D array on current processor */
    PetscInt Is, Js;
    PetscInt nx, ny;
    PetscErrorCode ierr;
    ierr = DMDAGetCorners(DA2d, &Is, &Js, PETSC_NULL, &nx, &ny, PETSC_NULL); CHKERRQ(ierr);
    PetscInt Ie = Is + nx;
    PetscInt Je = Js + ny;

    PetscScalar **rhs;
    ierr = DMDAVecGetArray(DA2d, b, &rhs); CHKERRQ(ierr);

    double (*rhs_func)(double) = NULL;
    if (boundary_condition == NEUMANN) {
        rhs_func = &cos;
    } else if (boundary_condition == DIRICHLET) {

        rhs_func = &sin;
    } /* else */

    /* Go over all the nodes within the current processor range, i.e.
  j = Js to Je-1 (y-direction)
   i = Is to Ie-1 (x-direction)
  and insert rhs values into the PETSc vector
 */
    int i, j;
    for (j=Js; j<Je; j++) {
        for (i=Is; i<Ie; i++) {

            PetscScalar x = i*dx + 0.5*dx;
            PetscScalar y = j*dy + 0.5*dy;

            PetscScalar coef = -(2.0*PI)*(2.0*PI) - (2.0*PI)*(2.0*PI);
            rhs[j][i] = -coef*rhs_func(2.0*PI*x)*rhs_func(2.0*PI*y);
            //PetscPrintf(PCW, "(i,j)=(%d,%d) (x,y)=(%f,%f) rhs:%f\n", i, j, x, y, rhs[j][i]);
        }
    }

    /* After working with the 2D array, always call this function */
    ierr = DMDAVecRestoreArray(DA2d, b, &rhs); CHKERRQ(ierr);
    //VecView(b, 0);

    return ierr;
}
/************************************************************************************************************************************************************/
/************************************************************************************************************************************************************/

/* This function sets up the analytical solution vector associated with the Poisson equation. 
This is done to validate our numerical results. 
Depending on the boudary condition, different analytical function is automatically chosen, i.e. 

1- Neumann BC, an_function   = cos(2pix)*cos(2piy) 
2- Dirichlet BC, an_function = sin(2pix)*sin(2piy) 
*/
PetscErrorCode setup_analytical_vector(Vec sol_an, DM DA2d, short int boundary_condition, PetscScalar dx, PetscScalar dy, PetscInt Nx, PetscInt Ny) {

    double (*rhs_func)(double) = NULL;
    if (boundary_condition == NEUMANN) {
        rhs_func = &cos;
    } else if (boundary_condition == DIRICHLET) {

        rhs_func = &sin;
    } else {
        PetscPrintf(PCW, "Error. Unknown boundary condition.\n");
    }

    /* Get the index range of the 2D array on current processor */
    PetscInt Is, Js;
    PetscInt nx, ny;
    PetscErrorCode ierr;
    ierr = DMDAGetCorners(DA2d, &Is, &Js, PETSC_NULL, &nx, &ny, PETSC_NULL); CHKERRQ(ierr);
    PetscInt Ie = Is + nx;
    PetscInt Je = Js + ny;

    PetscScalar **sol_;
    ierr = DMDAVecGetArray(DA2d, sol_an, &sol_); CHKERRQ(ierr);

    int i, j;
    for (j=Js; j<Je; j++) {
        for (i=Is; i<Ie; i++) {

            PetscScalar x = i*dx + 0.5*dx;
            PetscScalar y = j*dy + 0.5*dy;

            sol_[j][i] = rhs_func(2.0*PI*x)*rhs_func(2.0*PI*y);
        }
    }

    ierr = DMDAVecRestoreArray(DA2d, sol_an, &sol_); CHKERRQ(ierr);

    return ierr;
}
/************************************************************************************************************************************************************/
/************************************************************************************************************************************************************/

/* This function computes Norm1 and Norm2 errors of our numerical solution. */
PetscErrorCode compute_error(Vec sol, Vec sol_an, PetscScalar dx, PetscScalar dy) {

    PetscErrorCode ierr;

    /* First subtract solution from analytical solution. Result is stored in analytical solution vector. */
    ierr = VecAXPY(sol_an, -1.0, sol); CHKERRQ(ierr);

    PetscReal errors[2];

    /* Return norm1 and norm2 errors */
    ierr = VecNorm(sol_an, NORM_1_AND_2, errors); CHKERRQ(ierr);

    /* Multiply norm1 and norm2 with the grid spacing, for scaling purposed */
    PetscScalar E1 = errors[0]*dx*dy;
    PetscScalar E2 = errors[1]*sqrt(dx*dy);

    PetscPrintf(PCW, "Error analysis: \n");
    PetscPrintf(PCW, "   E1: %2.14lf\n", E1);
    PetscPrintf(PCW, "   E2: %2.14lf\n", E2);

    return ierr;
}