static const char help[] = "Solves obstacle problem in 2D as a variational inequality.\n\
Modified to include option of TAO complementarity solvers (SSILS/SSFLS/ASILS/ASFS). \n\
An elliptic problem with solution  u  constrained to be above a given function  psi. \n\
Exact solution is known.\n";

/*  Solve on a square R = {-2<x<2,-2<y<2}:
    u_{xx} + u_{yy} = 0
on the set where membrane is above obstacle.  Constraint is  u(x,y) >= psi(x,y).
Here psi is the upper hemisphere of the unit ball.  On the boundary of R
we have nonhomogenous Dirichlet boundary conditions coming from the exact solution.

Method is centered finite differences.

This example was contributed by Ed Bueler.  The exact solution is known for the
given psi and boundary values in question.  See
  https://github.com/bueler/fem-code-challenge/blob/master/obstacleDOC.pdf?raw=true.

Modified by Justin Chang to use TAO.

Example usage follows.

Get help:
  ./ex9_TAO -help

Basic run:
  ./ex9_TAO -solver_type <0/1>

  0 = SNES Variational Inequality
  1 = TAO Complementarity solvers

  Default SNESVI and TAO solvers are 'vinewtonrsls' and 'ssils' respectively.

*/

#include <petscdm.h>
#include <petscdmda.h>
#include <petscsnes.h>
#include <petsctao.h>

/* application context for obstacle problem solver */
typedef struct {
  DM da;
  Vec psi, uexact;
  Vec f; /* r - J*u if needed */
  Mat J; /* Jacobian J if needed */
} ObsCtx;
SNES snes;

extern PetscErrorCode FormPsiAndExactSoln(DM,ObsCtx*);
extern PetscErrorCode FormBounds(SNES,Vec,Vec);
extern PetscErrorCode FormFunctionLocal(DMDALocalInfo*,PetscScalar**,PetscScalar**,ObsCtx*);
extern PetscErrorCode FormJacobianLocal(DMDALocalInfo*,PetscScalar**,Mat,Mat,ObsCtx*);

/* Tao routines */
extern PetscErrorCode FormTaoBounds(Tao,Vec,Vec,void *);
extern PetscErrorCode FormTaoGradient(Tao,Vec,Vec,void *);
extern PetscErrorCode FormTaoJacobian(Tao,Vec,Mat,Mat,void *);

#undef __FUNCT__
#define __FUNCT__ "main"
int main(int argc,char **argv)
{
  PetscErrorCode      ierr;
  ObsCtx              user;
  Tao                 tao;
  Vec                 u;     /* solution */
  Vec                 xl,xu; /* lower and upper bounds */
  Vec                 r;     /* Residual, if needed */
  Vec                 Ju;    /* J*u vector, if needed */
  Vec                 c;     /* Constraints/gradient, if needed */
  DMDALocalInfo       info;
  PetscReal           error1,errorinf;
  PetscBool           flg;
  PetscInt            solver_type;

  /* Initialize */
  ierr = PetscInitialize(&argc,&argv,(char*)0,help);if (ierr) return ierr;
  solver_type = 0; /* Default to SNESVI */
  ierr = PetscOptionsGetInt(NULL,NULL, "-solver_type", &solver_type, &flg);CHKERRQ(ierr);
  ierr = DMDACreate2d(PETSC_COMM_WORLD,
                      DM_BOUNDARY_NONE, DM_BOUNDARY_NONE,
                      DMDA_STENCIL_STAR,
                      -11,-11,                   /* default to 10x10 grid */
                      PETSC_DECIDE,PETSC_DECIDE, /* number of processors in each dimension */
                      1,                         /* dof = 1 */
                      1,                         /* s = 1; stencil extends out one cell */
                      NULL,NULL,                 /* do not specify processor decomposition */
                      &user.da);CHKERRQ(ierr);

  ierr = DMCreateGlobalVector(user.da,&u);CHKERRQ(ierr);
  ierr = VecDuplicate(u,&(user.uexact));CHKERRQ(ierr);
  ierr = VecDuplicate(u,&(user.psi));CHKERRQ(ierr);
  ierr = VecDuplicate(u,&xl);CHKERRQ(ierr);
  ierr = VecDuplicate(u,&xu);CHKERRQ(ierr);

  ierr = DMDASetUniformCoordinates(user.da,-2.0,2.0,-2.0,2.0,0.0,1.0);CHKERRQ(ierr);

  ierr = FormPsiAndExactSoln(user.da,&user);CHKERRQ(ierr);
  ierr = VecSet(u,0.0);CHKERRQ(ierr);
  ierr = VecCopy(user.psi,xl);CHKERRQ(ierr);  /* u >= psi */
  ierr = VecSet(xu,PETSC_INFINITY);CHKERRQ(ierr);
  
  ierr = SNESCreate(PETSC_COMM_WORLD,&snes);CHKERRQ(ierr);
  ierr = SNESSetDM(snes,user.da);CHKERRQ(ierr);
  ierr = DMDASNESSetFunctionLocal(user.da,INSERT_VALUES,(PetscErrorCode (*)(DMDALocalInfo*,void*,void*,void*))FormFunctionLocal,&user);CHKERRQ(ierr);
  ierr = DMDASNESSetJacobianLocal(user.da,(PetscErrorCode (*)(DMDALocalInfo*,void*,Mat,Mat,void*))FormJacobianLocal,&user);CHKERRQ(ierr);

  /* Solver options */
  if (!solver_type) {
    ierr = SNESSetType(snes,SNESVINEWTONRSLS);CHKERRQ(ierr);
    ierr = SNESVISetVariableBounds(snes,xl,xu);CHKERRQ(ierr);
  } else {   
    ierr = DMCreateMatrix(user.da,&user.J);CHKERRQ(ierr);
    ierr = VecDuplicate(u, &c);CHKERRQ(ierr);
    ierr = VecDuplicate(u, &r);CHKERRQ(ierr);
    ierr = VecDuplicate(u, &Ju);CHKERRQ(ierr);
    ierr = VecDuplicate(u, &user.f);CHKERRQ(ierr);
    ierr = TaoCreate(PETSC_COMM_WORLD,&tao);CHKERRQ(ierr);
    ierr = TaoSetType(tao,TAOSSILS);CHKERRQ(ierr);
    ierr = TaoSetConstraintsRoutine(tao,c,FormTaoGradient,(void *)&user);CHKERRQ(ierr);
    ierr = TaoSetJacobianRoutine(tao,user.J,user.J,FormTaoJacobian,(void *)&user);CHKERRQ(ierr);
    ierr = TaoSetVariableBounds(tao,xl,xu);CHKERRQ(ierr);
    ierr = TaoSetInitialVector(tao,u);CHKERRQ(ierr);
    ierr = TaoSetFromOptions(tao);CHKERRQ(ierr);
  }
  ierr = SNESSetFromOptions(snes);CHKERRQ(ierr);

  /* report on setup */
  ierr = DMDAGetLocalInfo(user.da,&info); CHKERRQ(ierr);
  ierr = PetscPrintf(PETSC_COMM_WORLD,"setup done: grid  Mx,My = %D,%D  with spacing  dx,dy = %.4f,%.4f\n",
                     info.mx,info.my,(double)(4.0/(PetscReal)(info.mx-1)),(double)(4.0/(PetscReal)(info.my-1)));CHKERRQ(ierr);

  /* solve system */
  if (!solver_type) {
    /* VI */
    ierr = SNESSolve(snes,NULL,u);CHKERRQ(ierr);
  } else {
    /* TAO */
//    ierr = SNESComputeFunction(snes,u,r);CHKERRQ(ierr); /* Assemble residual */
//    ierr = SNESComputeJacobian(snes,u,user.J,user.J);CHKERRQ(ierr); /* Assemble Jacobian */
//    ierr = MatMult(user.J,u,Ju);CHKERRQ(ierr); /* Store J*u */
//    ierr = VecWAXPY(user.f,-1.0,Ju,r);CHKERRQ(ierr); /* Store f = r - J*u */
    ierr = TaoSolve(tao);CHKERRQ(ierr);
  }
  /* compare to exact */
  ierr = VecAXPY(u,-1.0,user.uexact);CHKERRQ(ierr); /* u <- u - uexact */
  ierr = VecNorm(u,NORM_1,&error1);CHKERRQ(ierr);
  error1 /= (PetscReal)info.mx * (PetscReal)info.my;
  ierr = VecNorm(u,NORM_INFINITY,&errorinf);CHKERRQ(ierr);
  ierr = PetscPrintf(PETSC_COMM_WORLD,"errors:     av |u-uexact|  = %.3e    |u-uexact|_inf = %.3e\n",(double)error1,(double)errorinf);CHKERRQ(ierr);

  /* Free work space.  */
  ierr = VecDestroy(&u);CHKERRQ(ierr);
  ierr = VecDestroy(&(user.psi));CHKERRQ(ierr);
  ierr = VecDestroy(&(user.uexact));CHKERRQ(ierr);
  if (solver_type == 1) {
    ierr = MatDestroy(&user.J);CHKERRQ(ierr);
    ierr = VecDestroy(&c);CHKERRQ(ierr);
    ierr = VecDestroy(&r);CHKERRQ(ierr);
    ierr = VecDestroy(&Ju);CHKERRQ(ierr);
    ierr = VecDestroy(&user.f);CHKERRQ(ierr);
    ierr = TaoDestroy(&tao);CHKERRQ(ierr);
  }
  ierr = SNESDestroy(&snes);CHKERRQ(ierr);
  ierr = DMDestroy(&user.da);CHKERRQ(ierr);
  ierr = PetscFinalize();CHKERRQ(ierr);
  return 0;
}


#undef __FUNCT__
#define __FUNCT__ "FormPsiAndExactSoln"
PetscErrorCode FormPsiAndExactSoln(DM da, ObsCtx *user) {
  PetscErrorCode ierr;
  DMDALocalInfo  info;
  PetscInt       i,j;
  DM             coordDA;
  Vec            coordinates;
  DMDACoor2d     **coords;
  PetscReal      **psi, **uexact, r;
  const PetscReal afree = 0.69797, A = 0.68026, B = 0.47152;

  PetscFunctionBeginUser;
  ierr = DMDAGetLocalInfo(da,&info); CHKERRQ(ierr);

  ierr = DMGetCoordinateDM(da, &coordDA);CHKERRQ(ierr);
  ierr = DMGetCoordinates(da, &coordinates);CHKERRQ(ierr);

  ierr = DMDAVecGetArray(coordDA, coordinates, &coords);CHKERRQ(ierr);
  ierr = DMDAVecGetArray(da, user->psi, &psi);CHKERRQ(ierr);
  ierr = DMDAVecGetArray(da, user->uexact, &uexact);CHKERRQ(ierr);
  for (j=info.ys; j<info.ys+info.ym; j++) {
    for (i=info.xs; i<info.xs+info.xm; i++) {
      r = PetscSqrtReal(PetscPowScalarInt(coords[j][i].x,2) + PetscPowScalarInt(coords[j][i].y,2));
      if (r <= 1.0) psi[j][i] = PetscSqrtReal(1.0 - r * r);
      else psi[j][i] = -1.0;
      if (r <= afree) uexact[j][i] = psi[j][i];  /* on the obstacle */
      else uexact[j][i] = - A * PetscLogReal(r) + B;   /* solves the laplace eqn */
    }
  }
  ierr = DMDAVecRestoreArray(da, user->psi, &psi);CHKERRQ(ierr);
  ierr = DMDAVecRestoreArray(da, user->uexact, &uexact);CHKERRQ(ierr);
  ierr = DMDAVecRestoreArray(coordDA, coordinates, &coords);CHKERRQ(ierr);
  PetscFunctionReturn(0);
}

#undef __FUNCT__
#define __FUNCT__ "FormTaoGradient"
/* Forms J*u + f where f = r - J*u */
PetscErrorCode FormTaoGradient(Tao tao, Vec X, Vec F, void *ctx)
{
  PetscErrorCode ierr;
  ObsCtx *user = (ObsCtx *)ctx;
  PetscInt i,j,mx,my,xs,xm,ys,ym;
  PetscReal dx,dy,uxx,uyy;
  PetscScalar **f, **x, **uexact;
  Vec localX;

  PetscFunctionBegin;
  ierr = DMDAGetInfo(user->da,PETSC_IGNORE,&mx,&my,PETSC_IGNORE,PETSC_IGNORE,PETSC_IGNORE,PETSC_IGNORE,PETSC_IGNORE,PETSC_IGNORE,PETSC_IGNORE,PETSC_IGNORE,PETSC_IGNORE,PETSC_IGNORE);CHKERRQ(ierr);
  ierr = DMDAGetCorners(user->da,&xs,&ys,NULL,&xm,&ym,NULL);CHKERRQ(ierr);
  //ierr = MatMult(user->J,x,g);CHKERRQ(ierr);
  //ierr = VecAXPY(g,1.0,user->f);CHKERRQ(ierr);
  
  ierr = DMGetLocalVector(user->da, &localX);CHKERRQ(ierr);
  ierr = DMGlobalToLocalBegin(user->da,X,INSERT_VALUES,localX);CHKERRQ(ierr);
  ierr = DMGlobalToLocalEnd(user->da,X,INSERT_VALUES,localX);CHKERRQ(ierr);
  ierr = DMDAVecGetArray(user->da, localX, &x);CHKERRQ(ierr);
  ierr = DMDAVecGetArray(user->da, F, &f);CHKERRQ(ierr);
  ierr = DMDAVecGetArray(user->da, user->uexact, &uexact);CHKERRQ(ierr);
  
  dx  = 4.0 / (PetscReal)(mx-1);
  dy  = 4.0 / (PetscReal)(my-1);
  for (j=ys; j<ys+ym; j++) {
    for (i=xs; i<xs+xm; i++) {
      if (i == 0 || j == 0 || i == mx-1 || j == my-1) {
        f[j][i] = 4.0*(x[j][i] - uexact[j][i]);
      } else {
        uxx     = dy*(x[j][i-1] - 2.0 * x[j][i] + x[j][i+1]) / dx;
        uyy     = dx*(x[j-1][i] - 2.0 * x[j][i] + x[j+1][i]) / dy;
        f[j][i] = -uxx - uyy;
      }
    }
  }
  ierr = DMDAVecRestoreArray(user->da, user->uexact, &uexact);CHKERRQ(ierr);
  ierr = DMDAVecRestoreArray(user->da, F, &f);CHKERRQ(ierr);
  ierr = DMDAVecRestoreArray(user->da, localX, &x);CHKERRQ(ierr);
  ierr = DMRestoreLocalVector(user->da,&localX);CHKERRQ(ierr);
  PetscFunctionReturn(0);
}

#undef __FUNCT__
#define __FUNCT__ "FormTaoJacobian"
/* Uses Jacobian J, do nothing else */
PetscErrorCode FormTaoJacobian(Tao tao, Vec x, Mat H, Mat Hpre, void *ctx)
{
  ObsCtx *user = (ObsCtx *)ctx;
  PetscErrorCode ierr;
  PetscInt i,j,mx,my,xs,xm,ys,ym;
  MatStencil col[5],row;
  PetscReal v[5],dx,dy,oxx,oyy;
  PetscBool assembled;
  
  PetscFunctionBegin;
  ierr = DMDAGetInfo(user->da,PETSC_IGNORE,&mx,&my,PETSC_IGNORE,PETSC_IGNORE,PETSC_IGNORE,PETSC_IGNORE,PETSC_IGNORE,PETSC_IGNORE,PETSC_IGNORE,PETSC_IGNORE,PETSC_IGNORE,PETSC_IGNORE);CHKERRQ(ierr);
  ierr = DMDAGetCorners(user->da,&xs,&ys,NULL,&xm,&ym,NULL);CHKERRQ(ierr);
 
  ierr = MatSetOption(H,MAT_IGNORE_OFF_PROC_ENTRIES,PETSC_TRUE);CHKERRQ(ierr);
  ierr = MatAssembled(H,&assembled);CHKERRQ(ierr);
  if (assembled) {ierr = MatZeroEntries(H);CHKERRQ(ierr);}

  dx  = 4.0 / (PetscReal)(mx-1);
  dy  = 4.0 / (PetscReal)(my-1);
  oxx = dy / dx;
  oyy = dx / dy;

  for (j=ys; j<ys+ym; j++) {
    for (i=xs; i<xs+xm; i++) {
      row.j = j; row.i = i;
      if (i == 0 || j == 0 || i == mx-1 || j == my-1) {
        v[0] = 4.0;
        ierr = MatSetValuesStencil(H,1,&row,1,&row,v,INSERT_VALUES);CHKERRQ(ierr);
      } else { 
        v[0] = -oyy;                 col[0].j = j - 1;  col[0].i = i;
        v[1] = -oxx;                 col[1].j = j;      col[1].i = i - 1;
        v[2] = 2.0 * (oxx + oyy);    col[2].j = j;      col[2].i = i;
        v[3] = -oxx;                 col[3].j = j;      col[3].i = i + 1;
        v[4] = -oyy;                 col[4].j = j + 1;  col[4].i = i;
        ierr = MatSetValuesStencil(H,1,&row,5,col,v,INSERT_VALUES);CHKERRQ(ierr);
      }
    }
  }

  ierr = MatAssemblyBegin(H,MAT_FINAL_ASSEMBLY);CHKERRQ(ierr);
  ierr = MatAssemblyEnd(H,MAT_FINAL_ASSEMBLY);CHKERRQ(ierr);
  ierr = PetscLogFlops(2.0*ym*xm);CHKERRQ(ierr);
  
  PetscFunctionReturn(0);
}

#undef __FUNCT__
#define __FUNCT__ "FormFunctionLocal"
/* FormFunctionLocal - Evaluates nonlinear function, F(x) on local process patch */
PetscErrorCode FormFunctionLocal(DMDALocalInfo *info,PetscScalar **x,PetscScalar **f,ObsCtx *user) {
  PetscErrorCode ierr;
  PetscInt       i,j;
  PetscReal      dx,dy,uxx,uyy;
  PetscReal      **uexact;  /* used for boundary values only */

  PetscFunctionBeginUser;
  dx = 4.0 / (PetscReal)(info->mx-1);
  dy = 4.0 / (PetscReal)(info->my-1);

  ierr = DMDAVecGetArray(info->da, user->uexact, &uexact);CHKERRQ(ierr);
  for (j=info->ys; j<info->ys+info->ym; j++) {
    for (i=info->xs; i<info->xs+info->xm; i++) {
      if (i == 0 || j == 0 || i == info->mx-1 || j == info->my-1) {
        f[j][i] = 4.0*(x[j][i] - uexact[j][i]);
      } else {
        uxx     = dy*(x[j][i-1] - 2.0 * x[j][i] + x[j][i+1]) / dx;
        uyy     = dx*(x[j-1][i] - 2.0 * x[j][i] + x[j+1][i]) / dy;
        f[j][i] = -uxx - uyy;
      }
    }
  }
  ierr = DMDAVecRestoreArray(info->da, user->uexact, &uexact);CHKERRQ(ierr);

  ierr = PetscLogFlops(10.0*info->ym*info->xm);CHKERRQ(ierr);
  PetscFunctionReturn(0);
}


#undef __FUNCT__
#define __FUNCT__ "FormJacobianLocal"
/* FormJacobianLocal - Evaluates Jacobian matrix on local process patch */
PetscErrorCode FormJacobianLocal(DMDALocalInfo *info,PetscScalar **x,Mat A,Mat jac, ObsCtx *user)
{
  PetscErrorCode ierr;
  PetscInt       i,j;
  MatStencil     col[5],row;
  PetscReal      v[5],dx,dy,oxx,oyy;

  PetscFunctionBeginUser;
  dx  = 4.0 / (PetscReal)(info->mx-1);
  dy  = 4.0 / (PetscReal)(info->my-1);
  oxx = dy / dx;
  oyy = dx / dy;

  for (j=info->ys; j<info->ys+info->ym; j++) {
    for (i=info->xs; i<info->xs+info->xm; i++) {
      row.j = j; row.i = i;
      if (i == 0 || j == 0 || i == info->mx-1 || j == info->my-1) { /* boundary */
        v[0] = 4.0;
        ierr = MatSetValuesStencil(jac,1,&row,1,&row,v,INSERT_VALUES);CHKERRQ(ierr);
      } else { /* interior grid points */
        v[0] = -oyy;                 col[0].j = j - 1;  col[0].i = i;
        v[1] = -oxx;                 col[1].j = j;      col[1].i = i - 1;
        v[2] = 2.0 * (oxx + oyy);    col[2].j = j;      col[2].i = i;
        v[3] = -oxx;                 col[3].j = j;      col[3].i = i + 1;
        v[4] = -oyy;                 col[4].j = j + 1;  col[4].i = i;
        ierr = MatSetValuesStencil(jac,1,&row,5,col,v,INSERT_VALUES);CHKERRQ(ierr);
      }
    }
  }

  /* Assemble matrix, using the 2-step process: */
  ierr = MatAssemblyBegin(jac,MAT_FINAL_ASSEMBLY);CHKERRQ(ierr);
  ierr = MatAssemblyEnd(jac,MAT_FINAL_ASSEMBLY);CHKERRQ(ierr);
  if (A != jac) {
    ierr = MatAssemblyBegin(A,MAT_FINAL_ASSEMBLY);CHKERRQ(ierr);
    ierr = MatAssemblyEnd(A,MAT_FINAL_ASSEMBLY);CHKERRQ(ierr);
  }
  ierr = PetscLogFlops(2.0*info->ym*info->xm);CHKERRQ(ierr);
  PetscFunctionReturn(0);
}