static char help[] = "This example solves a linear system in parallel with KSP.  The matrix\n\
uses arbitrary order polynomials for finite 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> -p <order> : mesh size and polynomial order\n\n";

#include "petscksp.h"

/* Declare user-defined routines */
extern PetscReal      src(PetscReal,PetscReal);
extern PetscReal      ubdy(PetscReal,PetscReal);
extern PetscReal      polyBasisFunc(PetscInt,PetscInt,PetscReal*,PetscReal);
extern PetscReal      derivPolyBasisFunc(PetscInt,PetscInt,PetscReal*,PetscReal);
extern PetscErrorCode Form1DElementMass(PetscReal,PetscInt,double*,double*,PetscReal*);
extern PetscErrorCode Form1DElementStiffness(PetscReal,PetscInt,double*,double*,PetscReal*);
extern PetscErrorCode Form2DElementMass(PetscInt,PetscReal*,PetscReal*);
extern PetscErrorCode Form2DElementStiffness(PetscInt,PetscReal*,PetscReal*,PetscReal*);
extern PetscErrorCode FormNodalRhs(PetscInt,PetscReal,PetscReal,PetscReal,double*,PetscReal*);
extern PetscErrorCode FormNodalSoln(PetscInt,PetscReal,PetscReal,PetscReal,double*,PetscReal*);
extern void leggaulob(double, double, double [], double [], int);
extern void qAndLEvaluation(int, double, double*, double*, double*);

#undef __FUNCT__
#define __FUNCT__ "main"
int main(int argc,char **args)
{
  PetscInt       p = 2, m = 5; 
  PetscInt       num1Dnodes, num2Dnodes;
  PetscReal      *Ke1D,*Ke2D,*Me1D,*Me2D;
  PetscReal      *r,*ue,val;   
  Vec            u,ustar,b,q;
  Mat            A,Mass;  
  KSP            ksp;     
  PetscInt       M,N;
  PetscMPIInt    rank,size;
  PetscReal      x,y,h,norm;
  PetscInt       *idx,indx,count,*rows,i,j,k,start,end,its;
  PetscReal      *rowsx,*rowsy;
  PetscReal      *gllNode, *gllWgts;
  PetscErrorCode ierr;

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

  /* 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);
  start = rank*(M/size) + ((M%size) < rank ? (M%size) : rank);
  end   = start + M/size + ((M%size) > rank); 

  /* Create matrix  */
  ierr  = MatCreate(PETSC_COMM_WORLD,&Mass);CHKERRQ(ierr);
  ierr  = MatSetSizes(Mass,PETSC_DECIDE,PETSC_DECIDE,N,N);CHKERRQ(ierr);
  ierr  = MatSetFromOptions(Mass);CHKERRQ(ierr);
  start = rank*(M/size) + ((M%size) < rank ? (M%size) : rank);
  end   = start + M/size + ((M%size) > rank);  

  /* Allocate element stiffness matrices */ 
  num1Dnodes = (p+1); 
  num2Dnodes = num1Dnodes*num1Dnodes;
  ierr = PetscMalloc((num1Dnodes*num1Dnodes)*sizeof(PetscReal),&Me1D);CHKERRQ(ierr);
  ierr = PetscMalloc((num1Dnodes*num1Dnodes)*sizeof(PetscReal),&Ke1D);CHKERRQ(ierr);
  ierr = PetscMalloc((num2Dnodes*num2Dnodes)*sizeof(PetscReal),&Me2D);CHKERRQ(ierr);
  ierr = PetscMalloc((num2Dnodes*num2Dnodes)*sizeof(PetscReal),&Ke2D);CHKERRQ(ierr);
  ierr = PetscMalloc(num2Dnodes*sizeof(PetscInt),&idx);CHKERRQ(ierr);
  ierr = PetscMalloc(num2Dnodes*sizeof(PetscReal),&r);CHKERRQ(ierr);
  ierr = PetscMalloc(num2Dnodes*sizeof(PetscReal),&ue);CHKERRQ(ierr);

  /* Allocate quadrature and create stiffness matrices */ 
  ierr = PetscMalloc((p+1)*sizeof(PetscReal),&gllNode);CHKERRQ(ierr); 
  ierr = PetscMalloc((p+1)*sizeof(PetscReal),&gllWgts);CHKERRQ(ierr);
  leggaulob(0.0,1.0,gllNode,gllWgts,p); // Get GLL nodes and weights
  ierr = Form1DElementMass(h,p,gllNode,gllWgts,Me1D);CHKERRQ(ierr);
  ierr = Form1DElementStiffness(h,p,gllNode,gllWgts,Ke1D);CHKERRQ(ierr);
  ierr = Form2DElementMass(p,Me1D,Me2D);CHKERRQ(ierr);
  ierr = Form2DElementStiffness(p,Ke1D,Me1D,Ke2D);CHKERRQ(ierr);

  /* Assemble matrix */ 
  for (i=start; i<end; i++) {
     indx = 0;
     for(k=0;k<(p+1);++k) {
       for(j=0;j<(p+1);++j) {
	 idx[indx++] = p*(p*m+1)*(i/m) + p*(i % m) + k*(p*m+1) + j;
       }
     }
     ierr = MatSetValues(A,num2Dnodes,idx,num2Dnodes,idx,Ke2D,ADD_VALUES);CHKERRQ(ierr);
     ierr = MatSetValues(Mass,num2Dnodes,idx,num2Dnodes,idx,Me2D,ADD_VALUES);CHKERRQ(ierr);
  }
  ierr = MatAssemblyBegin(A,MAT_FINAL_ASSEMBLY);CHKERRQ(ierr);
  ierr = MatAssemblyEnd(A,MAT_FINAL_ASSEMBLY);CHKERRQ(ierr);
  ierr = MatAssemblyBegin(Mass,MAT_FINAL_ASSEMBLY);CHKERRQ(ierr);
  ierr = MatAssemblyEnd(Mass,MAT_FINAL_ASSEMBLY);CHKERRQ(ierr);

  ierr = PetscFree(Me1D);CHKERRQ(ierr);
  ierr = PetscFree(Ke1D);CHKERRQ(ierr);
  ierr = PetscFree(Me2D);CHKERRQ(ierr);
  ierr = PetscFree(Ke2D);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(u,&q);CHKERRQ(ierr);
  ierr = PetscObjectSetName((PetscObject)q,"Right hand side 2");CHKERRQ(ierr);
  ierr = VecDuplicate(b,&ustar);CHKERRQ(ierr);
  ierr = VecSet(u,0.0);CHKERRQ(ierr);
  ierr = VecSet(b,0.0);CHKERRQ(ierr); 
  ierr = VecSet(q,0.0);CHKERRQ(ierr);  

  /* Assemble nodal right-hand-side and soln vector  */
  for (i=start; i<end; i++) {
     x = h*(i % m); 
     y = h*(i/m);
     indx = 0;
     for(k=0;k<(p+1);++k) {
       for(j=0;j<(p+1);++j) {
	 idx[indx++] = p*(p*m+1)*(i/m) + p*(i % m) + k*(p*m+1) + j;
       }
     }
     ierr = FormNodalRhs(p,x,y,h,gllNode,r);CHKERRQ(ierr);
     ierr = FormNodalSoln(p,x,y,h,gllNode,ue);CHKERRQ(ierr);
     ierr = VecSetValues(q,num2Dnodes,idx,r,INSERT_VALUES);CHKERRQ(ierr);
     ierr = VecSetValues(ustar,num2Dnodes,idx,ue,INSERT_VALUES);CHKERRQ(ierr);
  }
  ierr = VecAssemblyBegin(q);CHKERRQ(ierr);
  ierr = VecAssemblyEnd(q);CHKERRQ(ierr);
  ierr = VecAssemblyBegin(ustar);CHKERRQ(ierr);
  ierr = VecAssemblyEnd(ustar);CHKERRQ(ierr);

  ierr = PetscFree(r);CHKERRQ(ierr); 
  ierr = PetscFree(ue);CHKERRQ(ierr); 

  /* Get FE right-hand side vector */ 
  ierr = MatMult(Mass,q,b);CHKERRQ(ierr);

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

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

  /* Solve linear system */
  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);
  ierr = KSPSolve(ksp,b,u);CHKERRQ(ierr);

  /* Check error */
  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 = PetscFree(gllNode);CHKERRQ(ierr);
  ierr = PetscFree(gllWgts);CHKERRQ(ierr);

  ierr = KSPDestroy(ksp);CHKERRQ(ierr); 
  ierr = VecDestroy(u); CHKERRQ(ierr); ierr = VecDestroy(b); CHKERRQ(ierr);
  ierr = MatDestroy(A); CHKERRQ(ierr); ierr = MatDestroy(Mass); CHKERRQ(ierr);

  ierr = PetscFinalize();CHKERRQ(ierr);
  return 0;
}

/* --------------------------------------------------------------------- */

#undef __FUNCT__
#define __FUNCT__ "Form1DElementMass"
   /* 1d element stiffness mass matrix  */
PetscErrorCode Form1DElementMass(PetscReal H,PetscInt P,double* gqn,
                                 double* gqw,PetscReal *Me1D)
{
  int i,j,k;
  int indx;
  PetscFunctionBegin;
  for(j=0;j<(P+1);++j) {
    for(i=0;i<(P+1);++i) {
      indx = j*(P+1)+i;
      Me1D[indx] = 0.0;
      for(k=0; k<(P+1);++k) {
	Me1D[indx] += H*gqw[k]*polyBasisFunc(P,i,gqn,gqn[k])*polyBasisFunc(P,j,gqn,gqn[k]);
      }
    }
  }
  PetscFunctionReturn(0);
}

/* --------------------------------------------------------------------- */

#undef __FUNCT__
#define __FUNCT__ "Form1DElementStiffness"
   /* 1d element stiffness matrix for derivative */
PetscErrorCode Form1DElementStiffness(PetscReal H,PetscInt P,double* gqn,
                                      double* gqw,PetscReal *Ke1D)
{
  int i,j,k;
  int indx;
  PetscFunctionBegin;
  for(j=0;j<(P+1);++j) {
    for(i=0;i<(P+1);++i) {
      indx = j*(P+1)+i;
      Ke1D[indx] = 0.0;
      for(k=0; k<(P+1);++k) {
	Ke1D[indx] += (1./H)*gqw[k]*derivPolyBasisFunc(P,i,gqn,gqn[k])*derivPolyBasisFunc(P,j,gqn,gqn[k]);
      }
    }
  }
  
  PetscFunctionReturn(0);
}

/* --------------------------------------------------------------------- */

#undef __FUNCT__
#define __FUNCT__ "Form2DElementMass"
   /* element mass matrix */
PetscErrorCode Form2DElementMass(PetscInt P,PetscReal *Me1D,PetscReal *Me2D)
{
  int i1,j1,i2,j2;
  int indx1,indx2,indx3;
  PetscFunctionBegin;;
  for(j2=0;j2<(P+1);++j2) {
    for(i2=0; i2<(P+1);++i2) {
      for(j1=0;j1<(P+1);++j1) {
	for(i1=0;i1<(P+1);++i1) {
	  indx1 = j1*(P+1)+i1;
	  indx2 = j2*(P+1)+i2;
	  indx3 = (j2*(P+1)+j1)*(P+1)*(P+1)+(i2*(P+1)+i1);
	  Me2D[indx3] = Me1D[indx1]*Me1D[indx2]; 
	}
      }
    }
  }
  PetscFunctionReturn(0);
}

/* --------------------------------------------------------------------- */

#undef __FUNCT__
#define __FUNCT__ "Form2DElementStiffness"
   /* element stiffness for Laplacian */
PetscErrorCode Form2DElementStiffness(PetscInt P,PetscReal *Ke1D,PetscReal *Me1D,
                                      PetscReal *Ke2D)
{
  int i1,j1,i2,j2;
  int indx1,indx2,indx3;
  PetscFunctionBegin;
  for(j2=0;j2<(P+1);++j2) {
    for(i2=0; i2<(P+1);++i2) {
      for(j1=0;j1<(P+1);++j1) {
	for(i1=0;i1<(P+1);++i1) {
	  indx1 = j1*(P+1)+i1;
	  indx2 = j2*(P+1)+i2;
	  indx3 = (j2*(P+1)+j1)*(P+1)*(P+1)+(i2*(P+1)+i1);
	  Ke2D[indx3] = Ke1D[indx1]*Me1D[indx2] + Me1D[indx1]*Ke1D[indx2]; 
	}
      }
    }
  }
  PetscFunctionReturn(0);
}

/* --------------------------------------------------------------------- */

#undef __FUNCT__
#define __FUNCT__ "FormNodalRhs"
PetscErrorCode FormNodalRhs(PetscInt P,PetscReal x,PetscReal y,PetscReal H,double* nds,PetscReal *r)
{
  int i,j,indx;
  PetscFunctionBegin;
  indx=0;
  for(j=0;j<(P+1);++j) {
    for(i=0;i<(P+1);++i) {
      r[indx] = src(x+H*nds[i],y+H*nds[j]);
      indx++;
    }
  }
  PetscFunctionReturn(0);
}

/* --------------------------------------------------------------------- */

#undef __FUNCT__
#define __FUNCT__ "FormNodalRhs"
PetscErrorCode FormNodalSoln(PetscInt P,PetscReal x,PetscReal y,PetscReal H,double* nds,PetscReal *u)
{
  int i,j,indx;
  PetscFunctionBegin;
  indx=0;
  for(j=0;j<(P+1);++j) {
    for(i=0;i<(P+1);++i) {
      u[indx] = ubdy(x+H*nds[i],y+H*nds[j]);
      indx++;
    }
  }
  PetscFunctionReturn(0);
}

/* --------------------------------------------------------------------- */

#undef __FUNCT__
#define __FUNCT__ "polyBasisFunc"
PetscReal polyBasisFunc(PetscInt order, PetscInt basis, PetscReal* xLocVal, PetscReal xval) {
  PetscReal denominator = 1.;
  PetscReal numerator = 1.;
  PetscInt i=0;
  for (i=0; i<(order+1); i++) {
    if (i!=basis) {
      numerator *= (xval-xLocVal[i]);
      denominator *= (xLocVal[basis]-xLocVal[i]);
    }
  }
  return (numerator/denominator); 
}

/* --------------------------------------------------------------------- */

#undef __FUNCT__
#define __FUNCT__ "derivPolyBasisFunc"
PetscReal derivPolyBasisFunc(PetscInt order, PetscInt basis, PetscReal* xLocVal, PetscReal xval) {
  PetscReal denominator;
  PetscReal numerator;
  PetscReal numtmp;
  PetscInt i=0,j=0;
  denominator=1.;
  for (i=0; i<(order+1); i++) {
    if (i!=basis) {
      denominator *= (xLocVal[basis]-xLocVal[i]);
    }
  }
  numerator = 0.;
  for(j=0;j<(order+1);++j) {
    if(j != basis) {
      numtmp = 1.;
      for(i=0;i<(order+1);++i) {
	if(i!=basis && j!=i) {
	  numtmp *= (xval-xLocVal[i]);
	}
      }
      numerator += numtmp;
    }
  }

  return (numerator/denominator); 
}
  
/* --------------------------------------------------------------------- */

PetscReal ubdy(PetscReal x,PetscReal y)
{
  return x*x*y*y;
}

PetscReal src(PetscReal x,PetscReal y)
{
  return -2.*y*y - 2.*x*x;
}
/* --------------------------------------------------------------------- */

void leggaulob(double x1, double x2, double x[], double w[], int n) 
/*******************************************************************************
Given the lower and upper limits of integration x1 and x2, and given n, this
routine returns arrays x[0..n-1] and w[0..n-1] of length n, containing the abscissas
and weights of the Gauss-Lobatto-Legendre n-point quadrature formula.
*******************************************************************************/
{
  int j,m;
  double z1,z,xm,xl,q,qp,Ln,scale;
  double pi=3.1415926535897932385;
  if(n==1) {
    x[0]=x1;   /* Scale the root to the desired interval, */
    x[1]=x2;   /* and put in its symmetric counterpart.   */
    w[0]=1.;   /* Compute the weight */
    w[1]=1.;   /* and its symmetric counterpart. */
  } else {
    x[0]=x1;   /* Scale the root to the desired interval, */
    x[n]=x2;   /* and put in its symmetric counterpart.   */
    w[0]=2./(n*(n+1));;   /* Compute the weight */
    w[n]=2./(n*(n+1));   /* and its symmetric counterpart. */
    m=(n+1)/2; /* The roots are symmetric, so we only find half of them. */
    xm=0.5*(x2+x1);
    xl=0.5*(x2-x1);
    for (j=1;j<=(m-1);j++) { /* Loop over the desired roots. */
      z=-1.0*cos((pi*(j+0.25)/(n))-(3.0/(8.0*n*pi))*(1.0/(j+0.25)));
      /* Starting with the above approximation to the ith root, we enter */
      /* the main loop of refinement by Newton's method.                 */
      do {
	qAndLEvaluation(n,z,&q,&qp,&Ln);
	z1=z;
	z=z1-q/qp; /* Newton's method. */
      } while (fabs(z-z1) > 3.0e-11);
      qAndLEvaluation(n,z,&q,&qp,&Ln);
      x[j]=xm+xl*z;      /* Scale the root to the desired interval, */
      x[n-j]=xm-xl*z;  /* and put in its symmetric counterpart.   */
      w[j]=2.0/(n*(n+1)*Ln*Ln); /* Compute the weight */
      w[n-j]=w[j];                 /* and its symmetric counterpart. */
    }
  }
  if(n%2==0) {
    qAndLEvaluation(n,0.0,&q,&qp,&Ln);
    x[n/2]=(x2-x1)/2.0;
    w[n/2]=2.0/(n*(n+1)*Ln*Ln);
  }
  /* scale the weights according to mapping from [-1,1] to [0,1] */
  scale = (x2-x1)/2.0;
  for(j=0;j<=n;++j) w[j] = w[j]*scale;
}


/******************************************************************************/
void qAndLEvaluation(int n, double x, double* q, double* qp, double* Ln)
/*******************************************************************************
Compute the polynomial qn(x) = L_{N+1}(x) - L_{n-1}(x) and its derivative in 
addition to L_N(x) as these are needed for the GLL points.  See the book titled
"Implementing Spectral Methods for Partial Differential Equations: Algorithms 
for Scientists and Engineers" by David A. Kopriva.
*******************************************************************************/
{
  int k;

  double Lnp;
  double Lnp1, Lnp1p;
  double Lnm1, Lnm1p;
  double Lnm2, Lnm2p; 

  Lnm1 = 1.0;
  *Ln = x;
  Lnm1p = 0.0;
  Lnp = 1.0;

  for(k=2; k<=n; ++k) {
    Lnm2 = Lnm1;
    Lnm1 = *Ln;
    Lnm2p = Lnm1p;
    Lnm1p = Lnp;
    *Ln = (2.*k-1.)/(1.0*k)*x*Lnm1 - (k-1.)/(1.0*k)*Lnm2;
    Lnp = Lnm2p + (2.0*k-1)*Lnm1;
  }
  k=n+1;
  Lnp1 = (2.*k-1.)/(1.0*k)*x*(*Ln) - (k-1.)/(1.0*k)*Lnm1;
  Lnp1p = Lnm1p + (2.0*k-1)*(*Ln);
  *q = Lnp1 - Lnm1;
  *qp = Lnp1p - Lnm1p;
}