static char help[] = "Solves a linear system in parallel with KSP.  The matrix\n\
uses simple bilinear elements on the unit square.  To test the parallel\n\
matrix assembly, the matrix is intentionally laid out across processors\n\
differently from the way it is assembled.  Input arguments are:\n\
  -m <size> : problem size\n\n";

/*T
   Concepts: KSP^basic parallel example
   Concepts: Matrices^inserting elements by blocks
   Processors: n
T*/

/* 
  Include "petscksp.h" so that we can use KSP solvers.  Note that this file
  automatically includes:
     petscsys.h       - base PETSc routines   petscvec.h - vectors
     petscmat.h - matrices
     petscis.h     - index sets            petscksp.h - Krylov subspace methods
     petscviewer.h - viewers               petscpc.h  - preconditioners
*/
#include "petscksp.h"

/* Declare user-defined routines */
extern PetscErrorCode FormElementStiffness(PetscReal,PetscScalar*);
extern PetscErrorCode FormElementRhs(PetscReal,PetscReal,PetscReal,PetscScalar*);
extern PetscErrorCode VecView_VTK(Vec, const char [], const PetscInt, const PetscInt);

#undef __FUNCT__
#define __FUNCT__ "main"
int main(int argc,char **args)
{
  Vec            u,b,ustar; /* approx solution, RHS, exact solution */
  Mat            A;           /* linear system matrix */
  KSP            ksp;         /* Krylov subspace method context */
  PetscInt       N;           /* dimension of system (global) */
  PetscInt       M;           /* number of elements (global) */
  PetscMPIInt    rank;        /* processor rank */
  PetscMPIInt    size;        /* size of communicator */
  PetscScalar    Ke[16];      /* element matrix */
  PetscScalar    r[4];        /* element vector */
  PetscReal      h;           /* mesh width */
  PetscReal      norm;        /* norm of solution error */
  PetscReal      x,y;
  PetscScalar    val;
  PetscErrorCode ierr;
  PetscInt       idx[4],count,*rows,i,m = 5,start,end,its;

  PetscInitialize(&argc,&args,(char *)0,help);
  ierr = PetscOptionsGetInt(PETSC_NULL,"-m",&m,PETSC_NULL);CHKERRQ(ierr);
  N = (m+1)*(m+1);
  M = m*m;
  h = 1.0/m;
  ierr = MPI_Comm_rank(PETSC_COMM_WORLD,&rank);CHKERRQ(ierr);
  ierr = MPI_Comm_size(PETSC_COMM_WORLD,&size);CHKERRQ(ierr);

  /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - 
         Compute the matrix and right-hand-side vector that define
         the linear system, Au = b.
     - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */

  /* 
     Create stiffness matrix
  */
  ierr = MatCreate(PETSC_COMM_WORLD,&A);CHKERRQ(ierr);
  ierr = MatSetSizes(A,PETSC_DECIDE,PETSC_DECIDE,N,N);CHKERRQ(ierr);
  ierr = MatSetFromOptions(A);CHKERRQ(ierr);
  //----------------------------------------------------------------------------
  // modification to preallocate by Jed in petsc-user email found by searching "ex3.c"
  ierr = MatSeqAIJSetPreallocation(A,9,PETSC_NULL);CHKERRQ(ierr);
  ierr = MatMPIAIJSetPreallocation(A,9,PETSC_NULL,5,PETSC_NULL);CHKERRQ(ierr); /* More than necessary */
  //----------------------------------------------------------------------------
  start = rank*(M/size) + ((M%size) < rank ? (M%size) : rank);
  end   = start + M/size + ((M%size) > rank); 

  /*
     Assemble matrix
  */
  ierr = FormElementStiffness(h*h,Ke);
  for (i=start; i<end; i++) {
     /* location of lower left corner of element */
     x = h*(i % m); y = h*(i/m); 
     /* node numbers for the four corners of element */
     idx[0] = (m+1)*(i/m) + (i % m);
     idx[1] = idx[0]+1; idx[2] = idx[1] + m + 1; idx[3] = idx[2] - 1;
     ierr = MatSetValues(A,4,idx,4,idx,Ke,ADD_VALUES);CHKERRQ(ierr);
  }
  ierr = MatAssemblyBegin(A,MAT_FINAL_ASSEMBLY);CHKERRQ(ierr);
  ierr = MatAssemblyEnd(A,MAT_FINAL_ASSEMBLY);CHKERRQ(ierr);

  /*
     Create right-hand-side and solution vectors
  */
  ierr = VecCreate(PETSC_COMM_WORLD,&u);CHKERRQ(ierr); 
  ierr = VecSetSizes(u,PETSC_DECIDE,N);CHKERRQ(ierr); 
  ierr = VecSetFromOptions(u);CHKERRQ(ierr);
  ierr = PetscObjectSetName((PetscObject)u,"Approx. Solution");CHKERRQ(ierr);
  ierr = VecDuplicate(u,&b);CHKERRQ(ierr);
  ierr = PetscObjectSetName((PetscObject)b,"Right hand side");CHKERRQ(ierr);
  ierr = VecDuplicate(b,&ustar);CHKERRQ(ierr);
  ierr = VecSet(u,0.0);CHKERRQ(ierr);
  ierr = VecSet(b,0.0);CHKERRQ(ierr);

  /* 
     Assemble right-hand-side vector
  */
  for (i=start; i<end; i++) {
     /* location of lower left corner of element */
     x = h*(i % m); y = h*(i/m); 
     /* node numbers for the four corners of element */
     idx[0] = (m+1)*(i/m) + (i % m);
     idx[1] = idx[0]+1; idx[2] = idx[1] + m + 1; idx[3] = idx[2] - 1;
     ierr = FormElementRhs(x,y,h*h,r);CHKERRQ(ierr);
     ierr = VecSetValues(b,4,idx,r,ADD_VALUES);CHKERRQ(ierr);
  }
  ierr = VecAssemblyBegin(b);CHKERRQ(ierr);
  ierr = VecAssemblyEnd(b);CHKERRQ(ierr);

  /* 
     Modify matrix and right-hand-side for Dirichlet boundary conditions
  */
  ierr = PetscMalloc(4*m*sizeof(PetscInt),&rows);CHKERRQ(ierr);
  for (i=0; i<m+1; i++) {
    rows[i] = i; /* bottom */
    rows[3*m - 1 +i] = m*(m+1) + i; /* top */
  }
  count = m+1; /* left side */
  for (i=m+1; i<m*(m+1); i+= m+1) {
    rows[count++] = i;
  }
  count = 2*m; /* left side */
  for (i=2*m+1; i<m*(m+1); i+= m+1) {
    rows[count++] = i;
  }
  for (i=0; i<4*m; i++) {
     x = h*(rows[i] % (m+1)); y = h*(rows[i]/(m+1)); 
     val = y;
     //ierr = VecSetValues(u,1,&rows[i],&val,INSERT_VALUES);CHKERRQ(ierr);
     ierr = VecSetValues(b,1,&rows[i],&val,INSERT_VALUES);CHKERRQ(ierr);
  }    
  ierr = MatZeroRows(A,4*m,rows,1.0);CHKERRQ(ierr);
  ierr = PetscFree(rows);CHKERRQ(ierr);

  ierr = VecAssemblyBegin(u);CHKERRQ(ierr);
  ierr = VecAssemblyEnd(u);CHKERRQ(ierr);
  ierr = VecAssemblyBegin(b);CHKERRQ(ierr); 
  ierr = VecAssemblyEnd(b);CHKERRQ(ierr);

  /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - 
                Create the linear solver and set various options
     - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */

  ierr = KSPCreate(PETSC_COMM_WORLD,&ksp);CHKERRQ(ierr);
  ierr = KSPSetOperators(ksp,A,A,DIFFERENT_NONZERO_PATTERN);CHKERRQ(ierr);
  //ierr = KSPSetInitialGuessNonzero(ksp,PETSC_TRUE);CHKERRQ(ierr);
  ierr = KSPSetFromOptions(ksp);CHKERRQ(ierr);

  /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - 
                      Solve the linear system
     - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */

  MPI_Barrier( PETSC_COMM_WORLD );
  const double time = MPI_Wtime();

  ierr = KSPSolve(ksp,b,u);CHKERRQ(ierr);

  MPI_Barrier( PETSC_COMM_WORLD );
  const double timeSolve = MPI_Wtime() - time;

  //----------------------------------------------------------------------------
  // write some useful data
  ierr = PetscPrintf(PETSC_COMM_WORLD,"Time taken for solve: %A \n", timeSolve);CHKERRQ(ierr);

  PetscInt           numIterations;
  KSPConvergedReason convergedReason;
  ierr = KSPGetIterationNumber( ksp, &numIterations ); CHKERRV( ierr ); 
  ierr = KSPGetConvergedReason( ksp, &convergedReason ); CHKERRV( ierr );
  ierr = PetscPrintf(PETSC_COMM_WORLD,"Iterations: %D\nConverged Reason: %D\n", numIterations, convergedReason);CHKERRQ(ierr);

  PCType  pcType;
  KSPType kspType;
  PC      precond;
  ierr = KSPGetType( ksp, &kspType ); CHKERRV( ierr ); 
  ierr = KSPGetPC( ksp, &precond ); CHKERRV( ierr ); 
  ierr = PCGetType( precond, &pcType ); CHKERRV( ierr );
  ierr = PetscPrintf(PETSC_COMM_WORLD,"ksp_type: %s\npc_type: %s\n", kspType, pcType);CHKERRQ(ierr);
  //----------------------------------------------------------------------------

  /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - 
                      Check solution and clean up
     - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */

  /* Check error */
  ierr = VecGetOwnershipRange(ustar,&start,&end);CHKERRQ(ierr);
  for (i=start; i<end; i++) {
     x = h*(i % (m+1)); y = h*(i/(m+1)); 
     val = y;
     ierr = VecSetValues(ustar,1,&i,&val,INSERT_VALUES);CHKERRQ(ierr);
  }
  ierr = VecAssemblyBegin(ustar);CHKERRQ(ierr);
  ierr = VecAssemblyEnd(ustar);CHKERRQ(ierr);

  // output solutions to VTK format
  ierr = VecOutputVTK(ustar, "ex3_exact.vtk", (m + 1), (m + 1)); CHKERRQ(ierr);
  ierr = VecOutputVTK(u, "ex3_numerical.vtk", (m + 1), (m + 1)); CHKERRQ(ierr);

  ierr = VecAXPY(u,-1.0,ustar);CHKERRQ(ierr);
  ierr = VecNorm(u,NORM_2,&norm);CHKERRQ(ierr);
  ierr = KSPGetIterationNumber(ksp,&its);CHKERRQ(ierr);
  //  ierr = PetscPrintf(PETSC_COMM_WORLD,"Norm of error %A Iterations %D\n",norm*h,its);CHKERRQ(ierr);
  ierr = PetscPrintf(PETSC_COMM_WORLD,"Norm of error: %A\n",norm);CHKERRQ(ierr);

  /* 
     Free work space.  All PETSc objects should be destroyed when they
     are no longer needed.
  */
  ierr = KSPDestroy(ksp);CHKERRQ(ierr); ierr = VecDestroy(u);CHKERRQ(ierr);
  ierr = VecDestroy(ustar);CHKERRQ(ierr); ierr = VecDestroy(b);CHKERRQ(ierr);
  ierr = MatDestroy(A);CHKERRQ(ierr);

  /*
     Always call PetscFinalize() before exiting a program.  This routine
       - finalizes the PETSc libraries as well as MPI
       - provides summary and diagnostic information if certain runtime
         options are chosen (e.g., -log_summary).
  */
  ierr = PetscFinalize();CHKERRQ(ierr);
  return 0;
}

/* --------------------------------------------------------------------- */
#undef __FUNCT__
#define __FUNCT__ "FormElementStiffness"
   /* element stiffness for Laplacian */
PetscErrorCode FormElementStiffness(PetscReal H,PetscScalar *Ke)
{
  PetscFunctionBegin;
  //----------------------------------------------------------------------------
  //! Looks like this might be wrong?
  // Ke[0]  = H/6.0;    Ke[1]  = -.125*H; Ke[2]  = H/12.0;   Ke[3]  = -.125*H;
  // Ke[4]  = -.125*H;  Ke[5]  = H/6.0;   Ke[6]  = -.125*H;  Ke[7]  = H/12.0;
  // Ke[8]  = H/12.0;   Ke[9]  = -.125*H; Ke[10] = H/6.0;    Ke[11] = -.125*H;
  // Ke[12] = -.125*H;  Ke[13] = H/12.0;  Ke[14] = -.125*H;  Ke[15] = H/6.0;
  //----------------------------------------------------------------------------
  Ke[ 0] =  2./3.; Ke[ 1] = -1./6.; Ke[ 2] = -1./3.; Ke[ 3] = -1./6.;
  Ke[ 4] = -1./6.; Ke[ 5] =  2./3.; Ke[ 6] = -1./6.; Ke[ 7] = -1./3.;
  Ke[ 8] = -1./3.; Ke[ 9] = -1./6.; Ke[10] =  2./3.; Ke[11] = -1./6.;
  Ke[12] = -1./6.; Ke[13] = -1./3.; Ke[14] = -1./6.; Ke[15] =  2./3.;
  //----------------------------------------------------------------------------
  PetscFunctionReturn(0);
}
/* --------------------------------------------------------------------- */
#undef __FUNCT__
#define __FUNCT__ "FormElementRhs"
PetscErrorCode FormElementRhs(PetscReal x,PetscReal y,PetscReal H,PetscScalar *r)
{
  PetscFunctionBegin;
  r[0] = 0.; r[1] = 0.; r[2] = 0.; r[3] = 0.0;
  PetscFunctionReturn(0);
}
/* --------------------------------------------------------------------- */
#undef __FUNCT__
#define __FUNCT__ "VecOutputVTK"
PetscErrorCode VecOutputVTK( Vec                     x,   // vector to plot
			     const char     filename[],   // basename for file
			     const PetscInt          m,   // number of grid points in x1-direction
			     const PetscInt          n  ) // number of grid points in x2-direction
{
  PetscFunctionBegin;

  // Petsc error handler
  PetscErrorCode ierr;

  // pre-conditions: x is of length m * n
  PetscInt s; 
  ierr = VecGetSize( x, &s); CHKERRQ(ierr); 
  if (s != m * n) SETERRQ(1, "SETERRQ: Size mismatch in VecOutputVTK");

  // MPI Communicator
  MPI_Comm comm;
  ierr = PetscObjectGetComm((PetscObject) x, &comm); CHKERRQ(ierr);

  // Petsc viewers
  PetscViewer viewer;

  // write VTK file
  ierr = PetscViewerASCIIOpen(comm, filename, &viewer); CHKERRQ(ierr);
  PetscScalar dx = (PetscScalar)(1. / ( m - 1 ));
  PetscScalar dy = (PetscScalar)(1. / ( n - 1 ));
  ierr = PetscViewerASCIISynchronizedPrintf(viewer, "# vtk DataFile Version 2.0\n" );   CHKERRQ(ierr);
  ierr = PetscViewerASCIISynchronizedPrintf(viewer, "Laplace Equation\n" );             CHKERRQ(ierr);
  ierr = PetscViewerASCIISynchronizedPrintf(viewer, "ASCII\n" );                        CHKERRQ(ierr);
  ierr = PetscViewerASCIISynchronizedPrintf(viewer, "DATASET STRUCTURED_POINTS\n" );    CHKERRQ(ierr);
  ierr = PetscViewerASCIISynchronizedPrintf(viewer, "DIMENSIONS %d %d %d\n", m, n, 1 ); CHKERRQ(ierr);
  ierr = PetscViewerASCIISynchronizedPrintf(viewer, "SPACING %f %f %d\n", dx, dy, 1 );  CHKERRQ(ierr);
  ierr = PetscViewerASCIISynchronizedPrintf(viewer, "ORIGIN 0 0 0\n" );                 CHKERRQ(ierr);
  ierr = PetscViewerASCIISynchronizedPrintf(viewer, "POINT_DATA %d\n", m*n );           CHKERRQ(ierr);
  ierr = PetscViewerASCIISynchronizedPrintf(viewer, "SCALARS scalars double 1\n" );     CHKERRQ(ierr);
  ierr = PetscViewerASCIISynchronizedPrintf(viewer, "LOOKUP_TABLE table\n" );           CHKERRQ(ierr);

  // write scalar data
  PetscScalar *array;
  PetscInt i, nlocal;
  ierr = VecGetLocalSize(x, &nlocal); CHKERRQ(ierr);
  ierr = VecGetArray(x, &array);      CHKERRQ(ierr);
  for(i = 0; i < nlocal; i++) {
      ierr = PetscViewerASCIISynchronizedPrintf(viewer, "%f\n", array[i]); CHKERRQ(ierr);
  }
  ierr = PetscViewerFlush(viewer);   CHKERRQ(ierr);
  ierr = VecRestoreArray(x, &array); CHKERRQ(ierr);
  
  // clean up
  ierr = PetscViewerFlush(viewer);   CHKERRQ(ierr);
  ierr = PetscViewerDestroy(viewer); CHKERRQ(ierr);

  PetscFunctionReturn(0);
}