#ifdef PETSC_RCS_HEADER
static char vcid[] = "$Id: gsles.c,v 1.5 2000/01/10 03:54:18 knepley Exp $";
#endif

#include "petscsles.h"            /*I "petscsles.h" I*/
#include "gsolver.h"              /*I "gsolver.h" I*/

#undef  __FUNCT__
#define __FUNCT__ "GVecSolutionKSPMonitor"
/*@
  GVecSolutionKSPMonitor - Monitors solution at each KSP iteration.

  Collective on KSP

  Input Parameters:
. ksp   - The Krylov subspace context
. it    - The number of iterations so far
. rnorm - The current (approximate) residual norm
. ctx   - A viewer

  Level: advanced

.keywords: grid vector, ksp, monitor, solution
.seealso: KSPDefaultMonitor(),KSPSetMonitor(),GVecResidualKSPMonitor(),
          GVecRhsKSPMonitor(), GVecErrorKSPMonitor()
@*/
int GVecSolutionKSPMonitor(KSP ksp, int it, PetscReal rnorm, void *ctx)
{
  PetscViewer viewer = (PetscViewer) ctx;
  Vec         x;
  int         ierr;

  PetscFunctionBegin;
  PetscValidHeaderSpecific(ksp,    KSP_COOKIE);
  PetscValidHeaderSpecific(viewer, PETSC_VIEWER_COOKIE);
  ierr = KSPBuildSolution(ksp, PETSC_NULL, &x);                                                           CHKERRQ(ierr);
  ierr = VecView(x, viewer);                                                                              CHKERRQ(ierr);
  PetscFunctionReturn(0);
}

#undef  __FUNCT__
#define __FUNCT__ "GVecResidualKSPMonitor"
/*@
  GVecResidualKSPMonitor - Monitors residual at each KSP iteration.

  Collective on KSP

  Input Parameters:
. ksp   - The Krylov subspace context
. it    - The number of iterations so far
. rnorm - The current (approximate) residual norm
. ctx   - A viewer

  Level: advanced

.keywords: grid vector, ksp, monitor, residual
.seealso: KSPDefaultMonitor(),KSPSetMonitor(),GVecSolutionKSPMonitor(),
          GVecRhsKSPMonitor(), GVecErrorKSPMonitor()
@*/
int GVecResidualKSPMonitor(KSP ksp, int it, PetscReal rnorm, void *ctx)
{
  PetscViewer viewer = (PetscViewer) ctx;
  Vec         r;
  int         ierr;

  PetscFunctionBegin;
  PetscValidHeaderSpecific(ksp,    KSP_COOKIE);
  PetscValidHeaderSpecific(viewer, PETSC_VIEWER_COOKIE);
  ierr = KSPBuildResidual(ksp, PETSC_NULL, PETSC_NULL, &r);                                               CHKERRQ(ierr);
  ierr = VecView(r, viewer);                                                                              CHKERRQ(ierr);
  PetscFunctionReturn(0);
}

#undef  __FUNCT__
#define __FUNCT__ "GVecRhsKSPMonitor"
/*@
  GVecRhsKSPMonitor - Displays the right hand side for the linear system at the first iteration.

  Collective on KSP

  Input Parameters:
. ksp   - The Krylov subspace context
. it    - The number of iterations so far
. rnorm - The current (approximate) residual norm
. ctx   - A viewer

  Level: advanced

.keywords: grid vector, ksp, monitor, rhs
.seealso: KSPDefaultMonitor(),KSPSetMonitor(),GVecSolutionKSPMonitor(),
          GVecResidualKSPMonitor(), GVecErrorKSPMonitor()
@*/
int GVecRhsKSPMonitor(KSP ksp, int it, PetscReal rnorm, void *ctx)
{
  PetscViewer viewer = (PetscViewer) ctx;
  Vec         b;
  int         ierr;

  PetscFunctionBegin;
  PetscValidHeaderSpecific(ksp,    KSP_COOKIE);
  PetscValidHeaderSpecific(viewer, PETSC_VIEWER_COOKIE);
  if (!it) {
    ierr = KSPGetRhs(ksp, &b);                                                                            CHKERRQ(ierr);
    ierr = VecView(b, viewer);                                                                            CHKERRQ(ierr);
  }
  PetscFunctionReturn(0);
}

#undef  __FUNCT__
#define __FUNCT__ "GVecErrorKSPMonitor"
/*@
  GVecErrorKSPMonitor - Displays the error at each iteration.

  Collective on KSP

  Input Parameters:
. ksp      - The Krylov subspace context
. it       - The number of iterations so far
. rnorm    - The current (approximate) residual norm
. errorCtx - A GVecErrorKSPMonitorCtx

  Notes: 
  The final argument to KSPSetMonitor() with this routine must be
  a pointer to a GVecErrorKSPMonitorCtx.

  Level: advanced

.keywords: grid vector, ksp, monitor, error
.seealso: KSPDefaultMonitor(),KSPSetMonitor(),,GVecSolutionKSPMonitor(),
          GVecResidualKSPMonitor(), GVecRhsKSPMonitor()
@*/
int GVecErrorKSPMonitor(KSP ksp, int it, PetscReal rnorm, void *errorCtx)
{
  GVecErrorKSPMonitorCtx *ctx  = (GVecErrorKSPMonitorCtx *) errorCtx;
  PetscScalar             mone = -1.0;
  GVec                    x, e;
  PetscReal               norm_2, norm_max;
  int                     ierr;

  PetscFunctionBegin;
  PetscValidHeaderSpecific(ksp, KSP_COOKIE);
  ierr = KSPBuildSolution(ksp, PETSC_NULL, &x);                                                           CHKERRQ(ierr);
  ierr = GVecGetWorkGVec(ctx->solution, &e);                                                              CHKERRQ(ierr);
  ierr = VecWAXPY(&mone, x, ctx->solution, e);                                                            CHKERRQ(ierr);
  ierr = VecView(e, ctx->error_viewer);                                                                   CHKERRQ(ierr);

  /* Compute 2-norm and max-norm of error */
  if (ctx->norm_error_viewer) {
    ierr = VecNorm(e, NORM_2, &norm_2);                                                                   CHKERRQ(ierr);
    ierr = VecNorm(e, NORM_MAX, &norm_max);                                                               CHKERRQ(ierr);
    if (ctx->isInner == PETSC_FALSE) {
      PetscViewerASCIIPrintf(ctx->norm_error_viewer, "Iteration %d residual norm %g error 2-norm %g error max-norm %g\n",
                             it, rnorm, norm_2, norm_max);
    } else {
      PetscViewerASCIIPrintf(ctx->norm_error_viewer, "  Inner iteration %d residual norm %g error 2-norm %g error max-norm %g\n",
                             it, rnorm, norm_2, norm_max);
    }
  }

  ierr = GVecRestoreWorkGVec(ctx->solution, &e);                                                          CHKERRQ(ierr);
  PetscFunctionReturn(0);
}

#include "petscmg.h"
#undef  __FUNCT__
#define __FUNCT__ "GVecPCMGSetMonitor"
/* ------------------------------------------------------------------------------*/
/*
  GVecPCMGSetMonitor - Sets GVec KSP monitors for each level of multigrid.

   Input Parameters:
.    pc - preconditioner context for multigrid
*/
int GVecPCMGSetMonitor(PC pc, int views)
{
  char         name[25];
  int          levels, i,ierr, solution,rhs,residual,all;
  SLES         subsles;
  KSP          ksp;
  PetscViewer *viewers;
  PetscTruth   match;

  PetscFunctionBegin;
  PetscValidHeaderSpecific(pc, PC_COOKIE);
  ierr = PetscTypeCompare((PetscObject) pc, PCMG, &match);                                                CHKERRQ(ierr);
  if (match == PETSC_FALSE) PetscFunctionReturn(0);

  solution = (views & GVECPCMGMONITOR_SOLUTION);
  residual = (views & GVECPCMGMONITOR_RESIDUAL);
  rhs      = (views & GVECPCMGMONITOR_RHS);
  all      = (solution != 0) + (residual != 0) + (rhs != 0);

  ierr = MGGetLevels(pc, &levels);                                                                        CHKERRQ(ierr);

  ierr = PetscMalloc(all*levels * sizeof(PetscViewer), &viewers);                                         CHKERRQ(ierr);

  for(i = 0; i < levels; i++) {
    ierr = MGGetSmoother(pc, i, &subsles);                                                                CHKERRQ(ierr);
    ierr = SLESGetKSP(subsles, &ksp);                                                                     CHKERRQ(ierr);
    if (rhs) {
      sprintf(name, "Right hand side %d",i);
      ierr = PetscViewerDrawOpen(PETSC_COMM_WORLD, 0, name, PETSC_DECIDE, PETSC_DECIDE, 200, 200, viewers); CHKERRQ(ierr);
      ierr = KSPSetMonitor(ksp, GVecRhsKSPMonitor, *viewers++, PETSC_NULL);                               CHKERRQ(ierr);
    }
    if (solution) {
      sprintf(name, "Solution %d",i);
      ierr = PetscViewerDrawOpen(PETSC_COMM_WORLD, 0, name, PETSC_DECIDE, PETSC_DECIDE, 200, 200, viewers); CHKERRQ(ierr);
      ierr = KSPSetMonitor(ksp, GVecSolutionKSPMonitor, *viewers++, PETSC_NULL);                          CHKERRQ(ierr);
    }
    if (residual) {
      sprintf(name, "Residual %d",i);
      ierr = PetscViewerDrawOpen(PETSC_COMM_WORLD, 0, name, PETSC_DECIDE, PETSC_DECIDE, 200, 200, viewers); CHKERRQ(ierr);
      ierr = KSPSetMonitor(ksp, GVecResidualKSPMonitor, *viewers++, PETSC_NULL);                          CHKERRQ(ierr);
    }
  }

  PetscFunctionReturn(0);
}

EXTERN_C_BEGIN
#undef  __FUNCT__
#define __FUNCT__ "GVecKSPOptionsChecker_Private"
int GVecKSPOptionsChecker_Private(KSP ksp)
{
  char       *prefix, string[64], *p;
  int         ierr, loc[4], nmax;
  MPI_Comm    comm;
  PetscViewer viewer;
  PetscDraw   draw;
  PetscTruth  opt;


  PetscFunctionBegin;
  PetscValidHeaderSpecific(ksp, KSP_COOKIE);
  ierr = KSPGetOptionsPrefix(ksp, &prefix);                                                               CHKERRQ(ierr);
  ierr = PetscObjectGetComm((PetscObject) ksp, &comm);                                                    CHKERRQ(ierr);

  nmax = 4;
  loc[0] = PETSC_DECIDE; loc[1] = PETSC_DECIDE; loc[2] = 300; loc[3] = 300;
  ierr = PetscOptionsGetIntArray(prefix, "-gvec_ksp_solutionmonitor", loc, &nmax, &opt);                  CHKERRQ(ierr);
  if (opt) {
    ierr = PetscViewerDrawOpen(comm, 0, 0, loc[0], loc[1], loc[2], loc[3], &viewer);                      CHKERRQ(ierr);
    ierr = PetscViewerDrawGetDraw(viewer, 0, &draw);                                                      CHKERRQ(ierr);
    ierr = PetscDrawSetTitle(draw, "Approx. Solution");                                                   CHKERRQ(ierr);
    PetscLogObjectParent(ksp, (PetscObject) viewer);
    ierr = KSPSetMonitor(ksp, GVecSolutionKSPMonitor, (void *) viewer, PETSC_NULL);                       CHKERRQ(ierr);
  }
  nmax = 4;
  loc[0] = PETSC_DECIDE; loc[1] = PETSC_DECIDE; loc[2] = 300; loc[3] = 300;
  ierr = PetscOptionsGetIntArray(prefix, "-gvec_ksp_residualmonitor", loc, &nmax, &opt);                  CHKERRQ(ierr);
  if (opt) {
    ierr = PetscViewerDrawOpen(comm, 0, 0, loc[0], loc[1], loc[2], loc[3], &viewer);                      CHKERRQ(ierr);
    ierr = PetscViewerDrawGetDraw(viewer, 0, &draw);                                                      CHKERRQ(ierr);
    ierr = PetscDrawSetTitle(draw, "Residual");                                                           CHKERRQ(ierr);
    PetscLogObjectParent(ksp, (PetscObject) viewer);
    ierr = KSPSetMonitor(ksp, GVecResidualKSPMonitor, (void *) viewer, PETSC_NULL);                       CHKERRQ(ierr);
  }
  nmax = 4;
  loc[0] = PETSC_DECIDE; loc[1] = PETSC_DECIDE; loc[2] = 300; loc[3] = 300;
  ierr = PetscOptionsGetIntArray(prefix, "-gvec_ksp_rhsmonitor", loc, &nmax, &opt);                       CHKERRQ(ierr);
  if (opt) {
    ierr = PetscViewerDrawOpen(comm, 0, 0, loc[0], loc[1], loc[2], loc[3], &viewer);                      CHKERRQ(ierr);
    ierr = PetscViewerDrawGetDraw(viewer, 0, &draw);                                                      CHKERRQ(ierr);
    ierr = PetscDrawSetTitle(draw, "Rhs");                                                                CHKERRQ(ierr);
    PetscLogObjectParent(ksp, (PetscObject) viewer);
    ierr = KSPSetMonitor(ksp, GVecRhsKSPMonitor, (void *) viewer, PETSC_NULL);                            CHKERRQ(ierr);
  }
  ierr = PetscOptionsGetString(prefix, "-gvec_mg_monitor", string, 64, &opt);                             CHKERRQ(ierr);
  if (opt) {
    PC  pc;
    int all = 0;

    ierr = PetscStrstr(string, "nosolution", &p);                                                         CHKERRQ(ierr);
    if (!p) all |= GVECPCMGMONITOR_SOLUTION;
    ierr = PetscStrstr(string, "noresidual", &p);                                                         CHKERRQ(ierr);
    if (!p) all |= GVECPCMGMONITOR_RESIDUAL;
    ierr = PetscStrstr(string, "norhs",      &p);                                                         CHKERRQ(ierr);
    if (!p) all |= GVECPCMGMONITOR_RHS;
    ierr = KSPGetPC(ksp, &pc);                                                                            CHKERRQ(ierr);
    ierr = GVecPCMGSetMonitor(pc, all);                                                                   CHKERRQ(ierr);
  }
  ierr = PetscOptionsHasName(PETSC_NULL, "-help", &opt);                                                  CHKERRQ(ierr);
  if (opt) {
    char *pprefix;
    int   len = 2;

    if (prefix != PETSC_NULL) {
      ierr = PetscStrlen(prefix, &len);                                                                   CHKERRQ(ierr);
    }
    ierr = PetscMalloc((len+2) * sizeof(char), &pprefix);                                                 CHKERRQ(ierr);
    PetscStrcpy(pprefix, "-");
    if (prefix != PETSC_NULL)
      PetscStrcat(pprefix, prefix);
    PetscPrintf(comm," Additional KSP Monitor options for grid vectors\n");
    PetscPrintf(comm,"   %sgvec_ksp_solutionmonitor\n", pprefix);
    PetscPrintf(comm,"   %sgvec_ksp_residualmonitor\n", pprefix);
    PetscPrintf(comm,"   %sgvec_ksp_rhsmonitor\n", pprefix);
    PetscPrintf(comm,"   %sgvec_mg_monitor [nosolution,norhs,noresidual]\n", pprefix);
    ierr = PetscFree(pprefix);                                                                            CHKERRQ(ierr);
  }
  PetscFunctionReturn(0);
}
EXTERN_C_END

/*--------------------------------------------------- Uzawa Methods --------------------------------------------------*/
#undef  __FUNCT__
#define __FUNCT__ "GMatCreateUzawa"
/*@
   GMatCreateUzawa - Creates a container matrix for the Uzawa system matrix B^T A^{-1} B.

   Input Parameter:
+  A    - The square (0,0) block of the original saddle-point matrix
.  B    - The rectangular (0,1) block of the original saddle-point matrix
-  sles - The solver used to invert A

   Output Parameter:
.  gmat - The discrete operator

  Level: advanced

.seealso: GVecCreate()
@*/
int GMatCreateUzawa(GMat A, GMat B, SLES sles, GMat *gmat)
{
  UzawaContext *uCtx;
  MPI_Comm      comm;
  Grid          grid, grid2;
  int           M, m, N, n;
  int           O, o, P, p;
  int           ierr;

  PetscFunctionBegin;
  PetscValidHeaderSpecific(A,    MAT_COOKIE);
  PetscValidHeaderSpecific(B,    MAT_COOKIE);
  PetscValidHeaderSpecific(sles, SLES_COOKIE);
  PetscValidPointer(gmat);
  ierr = GMatGetGrid(A, &grid);                                                                           CHKERRQ(ierr);
  ierr = GMatGetGrid(B, &grid2);                                                                          CHKERRQ(ierr);
  PetscValidHeaderSpecific(grid,  GRID_COOKIE);
  PetscValidHeaderSpecific(grid2, GRID_COOKIE);
  if (grid != grid2) SETERRQ(PETSC_ERR_ARG_INCOMP, "Matrices must all have the same underlying grid.");
  ierr = PetscNew(UzawaContext, &uCtx);                                                                   CHKERRQ(ierr);
  uCtx->A    = A;
  uCtx->B    = B;
  uCtx->sles = sles;
  ierr = MatGetSize(A, &M, &N);                                                                           CHKERRQ(ierr);
  ierr = MatGetLocalSize(A, &m, &n);                                                                      CHKERRQ(ierr);
  ierr = MatGetSize(B, &O, &P);                                                                           CHKERRQ(ierr);
  ierr = MatGetLocalSize(B, &o, &p);                                                                      CHKERRQ(ierr);
  if ((m != o) || (n != o) || (M != O) || (N != O)) {
    SETERRQ(PETSC_ERR_ARG_INCOMP, "Incommensurate sizes, B^T A B must be a legal matrix");
  }
  ierr = PetscObjectGetComm((PetscObject) grid, &comm);                                                   CHKERRQ(ierr);
  ierr = MatCreateShell(comm, p, p, P, P, uCtx, gmat);                                                    CHKERRQ(ierr);
  ierr = MatShellSetOperation(*gmat, MATOP_MULT, (void (*)(void)) GMatMatMultUzawa);                      CHKERRQ(ierr);
  ierr = VecCreateMPI(comm, m, M, &uCtx->work);                                                           CHKERRQ(ierr);
  ierr = VecDuplicate(uCtx->work, &uCtx->work2);                                                          CHKERRQ(ierr);
  PetscFunctionReturn(0);
}

#undef __FUNCT__  
#define __FUNCT__ "GMatMatMultUzawa"
/*@
  GMatMatMultUzawa - This function applies B^T A^{-1} B to a vector,
  which is the Schur complement matrix.

  Input Parameters:
+ mat - The grid matrix
- x   - The input grid vector

  Output Parameter:
. y - The output grid vector B^T A^{-1} B x

  Level: advanced

.seealso: GMatEvaluateOperatorGalerkin
@*/
int GMatMatMultUzawa(GMat mat, GVec x, GVec y)
{
  UzawaContext *ctx;
  int           its;
  int           ierr;

  PetscFunctionBegin;
  PetscValidHeaderSpecific(mat, MAT_COOKIE);
  PetscValidHeaderSpecific(x,   VEC_COOKIE);
  PetscValidHeaderSpecific(y,   VEC_COOKIE);
  ierr = MatShellGetContext(mat, (void **) &ctx);                                                         CHKERRQ(ierr);
  ierr = MatMult(ctx->B, x, ctx->work);                                                                   CHKERRQ(ierr);
  ierr = SLESSolve(ctx->sles, ctx->work, ctx->work2, &its);                                               CHKERRQ(ierr);
  ierr = MatMultTranspose(ctx->B, ctx->work2, y);                                                         CHKERRQ(ierr);
  PetscFunctionReturn(0);
}

#undef  __FUNCT__
#define __FUNCT__ "GMatDestroyUzawa"
/*@
   GMatDestroyUzawa - Destroys a container matrix for the Uzawa
   system matrix B^T A B.

   Input Parameter:
.  gmat - The container matrix from GMatCreateUzawa()

  Level: advanced

.seealso: GVecCreate()
@*/
int GMatDestroyUzawa(GMat gmat)
{
  UzawaContext *ctx;
  int           ierr;

  PetscFunctionBegin;
  PetscValidHeaderSpecific(gmat, MAT_COOKIE);

  ierr = MatShellGetContext(gmat, (void **) &ctx);                                                        CHKERRQ(ierr);
  if (ctx != PETSC_NULL) {
    ierr = VecDestroy(ctx->work);                                                                         CHKERRQ(ierr);
    ierr = VecDestroy(ctx->work2);                                                                        CHKERRQ(ierr);
    ierr = PetscFree(ctx);                                                                                CHKERRQ(ierr);
  }
  ierr = MatDestroy(gmat);                                                                                CHKERRQ(ierr);
  PetscFunctionReturn(0);
}