#ifdef PETSC_RCS_HEADER
static char vcid[] = "$Id: jacobic.c,v 1.1 1999/05/20 21:56:24 knepley Exp $";
#endif
/*
   Defines a Jacobi preconditioner for any constrained Mat implementation
*/
#include "src/sles/pc/pcimpl.h"   /*I "petscpc.h" I*/
#include "gvec.h"

/* We are cheating here since the interface is deficient */
typedef struct {
  Vec        diag;
  Vec        diagsqrt;
  PetscTruth userowmax;
} PC_Jacobi;

typedef struct {
  PC jacobi;
} PC_Jacobi_Constrained;

#undef __FUNCT__  
#define __FUNCT__ "PCSetUp_Jacobi_Constrained"
static int PCSetUp_Jacobi_Constrained(PC pc) {
  PC_Jacobi_Constrained *jacConst     = (PC_Jacobi_Constrained *) pc->data;
  PC_Jacobi             *jac          = (PC_Jacobi *) jacConst->jacobi->data;
  PetscTruth             zeroDetected = PETSC_FALSE;
  Grid                   grid;
  Vec                    diag, diagsqrt;
  PetscScalar           *x;
  int                    i, n;
  int                    ierr;

  PetscFunctionBegin;
  ierr = PCSetVector(jacConst->jacobi, pc->vec);                                                          CHKERRQ(ierr);
  ierr = PCSetOperators(jacConst->jacobi, pc->mat, pc->pmat, pc->flag);                                   CHKERRQ(ierr);
  ierr = GMatGetGrid(pc->pmat, &grid);                                                                    CHKERRQ(ierr);
  if (grid == PETSC_NULL) SETERRQ(PETSC_ERR_SUP, "Constrained Jacobi is only for Grid matrices");

  jacConst->jacobi->setupcalled = 2;
  diag     = jac->diag;
  diagsqrt = jac->diagsqrt;
  if (jac->userowmax) SETERRQ(PETSC_ERR_SUP, "Not currently supported for constrained matrices");
  if (diag) {
    ierr = GMatGetDiagonalConstrained(pc->pmat, diag);                                                    CHKERRQ(ierr);
    ierr = VecReciprocal(diag);                                                                           CHKERRQ(ierr);
    ierr = VecGetLocalSize(diag, &n);                                                                     CHKERRQ(ierr);
    ierr = VecGetArray(diag, &x);                                                                         CHKERRQ(ierr);
    for(i = 0; i < n; i++) {
      if (x[i] == 0.0) {
        x[i] = 1.0;
        zeroDetected = PETSC_TRUE;
      }
    }
    ierr = VecRestoreArray(diag, &x);                                                                     CHKERRQ(ierr);
  }
  if (diagsqrt) {
    ierr = GMatGetDiagonalConstrained(pc->pmat, diagsqrt);                                                CHKERRQ(ierr);
    ierr = VecGetLocalSize(diagsqrt, &n);                                                                 CHKERRQ(ierr);
    ierr = VecGetArray(diagsqrt, &x);                                                                     CHKERRQ(ierr);
    for(i = 0; i < n; i++) {
      if (x[i] != 0.0) {
        x[i] = 1.0/sqrt(PetscAbsScalar(x[i]));
      } else {
        x[i] = 1.0;
        zeroDetected = PETSC_TRUE;
      }
    }
    ierr = VecRestoreArray(diagsqrt, &x);                                                                 CHKERRQ(ierr);
  }

  if (zeroDetected == PETSC_TRUE) {
    PetscLogInfo(pc, "PCSetUp_Jacobi_Constrained:WARNING: Zero detected in diagonal while building Jacobi preconditioner\n");
  }
  PetscFunctionReturn(0);
}

#undef __FUNCT__  
#define __FUNCT__ "PCDestroy_Jacobi_Constrained"
static int PCDestroy_Jacobi_Constrained(PC pc) {
  PC_Jacobi_Constrained *jac = (PC_Jacobi_Constrained *) pc->data;
  int                    ierr;

  PetscFunctionBegin;
  ierr = PCDestroy(jac->jacobi);CHKERRQ(ierr);
  ierr = PetscFree(jac);CHKERRQ(ierr);
  PetscFunctionReturn(0);
}

#undef __FUNCT__  
#define __FUNCT__ "PCApply_Jacobi_Constrained"
static int PCApply_Jacobi_Constrained(PC pc, Vec x, Vec y) {
  PC_Jacobi_Constrained *jac = (PC_Jacobi_Constrained *) pc->data;
  PC_Jacobi             *j   = (PC_Jacobi *) jac->jacobi->data;
  int                    ierr;

  if (j->diag == PETSC_NULL) {
    ierr = VecDuplicate(pc->vec, &j->diag);                                                               CHKERRQ(ierr);
    PetscLogObjectParent(pc, j->diag);
    ierr = PCSetUp_Jacobi_Constrained(pc);                                                                CHKERRQ(ierr);
  }
  return PCApply(jac->jacobi, x, y);
}

#undef __FUNCT__  
#define __FUNCT__ "PCSetFromOptions_Jacobi_Constrained"
static int PCSetFromOptions_Jacobi_Constrained(PC pc) {
  PC_Jacobi_Constrained *jac = (PC_Jacobi_Constrained *) pc->data;
  return PCSetFromOptions(jac->jacobi);
}

#undef __FUNCT__  
#define __FUNCT__ "PCApplySymmetricLeftOrRight_Jacobi_Constrained"
static int PCApplySymmetricLeftOrRight_Jacobi_Constrained(PC pc, Vec x, Vec y) {
  PC_Jacobi_Constrained *jac = (PC_Jacobi_Constrained *) pc->data;
  PC_Jacobi             *j   = (PC_Jacobi *) jac->jacobi->data;
  int                    ierr;

  if (j->diagsqrt == PETSC_NULL) {
    ierr = VecDuplicate(pc->vec, &j->diagsqrt);                                                           CHKERRQ(ierr);
    PetscLogObjectParent(pc, j->diagsqrt);
    ierr = PCSetUp_Jacobi_Constrained(pc);                                                                CHKERRQ(ierr);
  }
  return PCApplySymmetricLeft(jac->jacobi, x, y);
}

EXTERN_C_BEGIN
#undef __FUNCT__  
#define __FUNCT__ "PCCreate_Jacobi_Constrained"
int PCCreate_Jacobi_Constrained(PC pc) {
  PC_Jacobi_Constrained *jac;
  int                    ierr;

  PetscFunctionBegin;
  ierr = PetscNew(PC_Jacobi_Constrained, &jac);                                                           CHKERRQ(ierr);
  PetscLogObjectMemory(pc,sizeof(PC_Jacobi_Constrained));
  pc->data = (void*) jac;

  ierr = PCCreate(pc->comm, &jac->jacobi);                                                                CHKERRQ(ierr);
  ierr = PCSetOptionsPrefix(jac->jacobi, "jacobic");                                                      CHKERRQ(ierr);
  ierr = PCSetType(jac->jacobi, PCJACOBI);                                                                CHKERRQ(ierr);
  pc->ops->apply               = PCApply_Jacobi_Constrained;
  pc->ops->applytranspose      = PCApply_Jacobi_Constrained;
  pc->ops->setup               = PCSetUp_Jacobi_Constrained;
  pc->ops->destroy             = PCDestroy_Jacobi_Constrained;
  pc->ops->setfromoptions      = PCSetFromOptions_Jacobi_Constrained;
  pc->ops->view                = PETSC_NULL;
  pc->ops->applyrichardson     = PETSC_NULL;
  pc->ops->applysymmetricleft  = PCApplySymmetricLeftOrRight_Jacobi_Constrained;
  pc->ops->applysymmetricright = PCApplySymmetricLeftOrRight_Jacobi_Constrained;
  PetscFunctionReturn(0);
}
EXTERN_C_END