#include "tao_general.h"

#ifdef TAO_USE_PETSC

#include "taolinearsolver_petsc.h"

#include "src/ksp/ksp/kspimpl.h"
#include "src/ksp/ksp/impls/qcg/qcg.h"

#include "../matrix/taomat_petsc.h"
#include "taovec.h"
// #include "../vector/taovec_petsc.h"

int KSPDuplicate(KSP ksp1, KSP* kspout);
extern int PetscMapSame(PetscMap, PetscMap, PetscTruth *);

TaoLinearSolverPetsc::TaoLinearSolverPetsc(KSP S):TaoLinearSolver(){
  linear_its=0;
  rrmap=0;rcmap=0;
  ksp=0;
  this->pkspviewer=PETSC_VIEWER_STDOUT_WORLD;
  this->SetKSP(S);
  return;
}

TaoLinearSolverPetsc::~TaoLinearSolverPetsc(){
  if (ksp){
    KSPDestroy(ksp);
  }
  if (this->rrmap && this->rcmap){
    PetscObjectDereference((PetscObject)this->rrmap);
    PetscObjectDereference((PetscObject)this->rcmap);    
  };
  return;
}

#undef __FUNCT__
#define __FUNCT__ "TaoWrapKSP"
/*@C
   TaoWrapKSP - Create a new TaoLinearSolver object using PETSc KSP.

   Input Parameter:
.  S -  a KSP

   Output Parameter:
.  SS - new TaoMat

   Note:  
   A TaoLinearSolverPetsc is an object with the methods of an abstract
   TaoLinearSolver object.  A TaoLinearSolverPetsc contains an implementation 
   of the TaoLinearSolver methods.  Routines using these vectors should 
   declare a pointer to a TaoLinearSolver, assign this pointer to the address 
   of a TaoLinearSolver object, use the pointer to invoke methods on the 
   object, and use this pointer as an argument when calling other routines.  
   This usage is different from the usage of a PETSc KSP.  In PETSc, 
   applications will typically declare a KSP, and pass it as an argument into 
   routines.  That is, applications will typically declare a pointer to a 
   TaoLinearSolver and use the pointer, or declare a KSP and use it directly.

   Level: developer

@*/
int TaoWrapKSP( KSP S, TaoLinearSolverPetsc ** SS){
  PetscFunctionBegin;
  PetscValidHeaderSpecific(S,KSP_COOKIE,1);
  *SS = new  TaoLinearSolverPetsc(S);
  PetscFunctionReturn(0);
}

#undef __FUNCT__
#define __FUNCT__ "TaoLinearSolverGetKSP"
/*@C
   TaoLinearSolverGetKSP - If the TaoLinearSolver is of the TaoLinearSolverPetsc type, it gets the underlying PETSc KSP 

   Input Parameter:
+  TS - the TaoLinearSolver
-  S -  the address of KSP

   Output Parameter:
.  S -  address of the underlying PETSc KSP object


   Note:  
   This routine does not create a KSP.  It sets a pointer
   to the location of an existing KSP.

   Level: advanced

.seealso TaoVecGetPetscVec(), TaoMatGetPetscMat(), TaoGetKSP()
@*/
int TaoLinearSolverGetKSP( TaoLinearSolver *TS, KSP * S){
  PetscFunctionBegin;
  if (TS){
    *S=((TaoLinearSolverPetsc *)TS)->GetKSP();
    PetscValidHeaderSpecific(*S,KSP_COOKIE,2);
  } else {
    *S=0;
  }
  PetscFunctionReturn(0);
}

#undef __FUNCT__
#define __FUNCT__ "TaoLinearSolverPetsc::SetKSP"
int TaoLinearSolverPetsc::SetKSP(KSP ksp2){
  int info;
  PetscFunctionBegin;
  if (ksp2){
    PetscValidHeaderSpecific(ksp2,KSP_COOKIE,1);
    PetscObjectReference((PetscObject)ksp2); 
  }
  if (this->ksp){
    info=KSPDestroy(ksp); CHKERRQ(info);
  }
  this->ksp=ksp2;
  this->LinearSolverObject=(void*)ksp2; 
  PetscFunctionReturn(0);
}
#undef __FUNCT__
#define __FUNCT__ "TaoLinearSolverPetsc::GetKSP"
int TaoLinearSolverPetsc::GetKSP(KSP *ksp2){
  PetscFunctionBegin;
  if (ksp2){
    *ksp2=this->ksp;
  }
  PetscFunctionReturn(0);
}

#undef __FUNCT__
#define __FUNCT__ "TaoLinearSolverPetsc::PreSolve"
int TaoLinearSolverPetsc::PreSolve(TaoMat* m1){
  int info;
  TaoMatPetsc *MM=(TaoMatPetsc*)m1;
  Mat mm, mm_pre;
  MatStructure preflag;
  PetscMap rmap1,cmap1;
  PetscTruth samemap=PETSC_TRUE;
  KSP tksp, ksp1=this->ksp;

  PetscFunctionBegin;
  info = MM->GetMatrix(&mm,&mm_pre,&preflag);CHKERRQ(info);
  info = MatGetPetscMaps(mm,&rmap1,&cmap1);CHKERRQ(info);

  if (this->rrmap && this->rcmap){
    info=PetscMapSame(rmap1,this->rrmap, &samemap);CHKERRQ(info);
    if (samemap){
      info=PetscMapSame(cmap1,this->rcmap, &samemap);CHKERRQ(info);
    }
    PetscObjectDereference((PetscObject)this->rrmap);
    PetscObjectDereference((PetscObject)this->rcmap);    
  }

  if (samemap==PETSC_FALSE){
    info = KSPDuplicate(ksp1,&tksp); CHKERRQ(info);
    info = this->SetKSP(tksp); CHKERRQ(info);
    info = KSPDestroy(tksp);  CHKERRQ(info);
  }

  this->rrmap=rmap1; this->rcmap=cmap1;

  PetscObjectReference((PetscObject)rmap1);
  PetscObjectReference((PetscObject)cmap1);

  info = KSPSetOperators(this->ksp,mm,mm_pre,preflag);CHKERRQ(info);
  PetscFunctionReturn(0);
}

#undef __FUNCT__
#define __FUNCT__ "TaoLinearSolverPetsc::SetTrustRadius"
int TaoLinearSolverPetsc::SetTrustRadius(double rad){

  KSPType ktype;
  KSP_QCG *qcgP;
  int info;
  PetscTruth flg;

  PetscFunctionBegin;

  info=KSPGetType(ksp,&ktype); CHKERRQ(info);
  info = PetscStrcmp((char *)ktype,KSPQCG,&flg); CHKERRQ(info);
  if (flg==PETSC_TRUE){
    qcgP = (KSP_QCG*) ksp->data;
    qcgP->delta=rad;
  }
  PetscFunctionReturn(0);
}

#undef __FUNCT__
#define __FUNCT__ "TaoLinearSolverPetsc::Solve"
int TaoLinearSolverPetsc::Solve(TaoVec* v1, TaoVec* v2, TaoTruth *flag){
  Vec vv=(Vec) v1->VecObject;
  Vec ww=(Vec) v2->VecObject;
  int info,its;

  PetscFunctionBegin;
  info = KSPSolve(ksp,vv,ww); CHKERRQ(info);
  info = KSPGetIterationNumber(ksp, &its); CHKERRQ(info);

  this->linear_its=PetscMax(its,-its);
  if (its>0) *flag=TAO_TRUE; 
  else *flag=TAO_FALSE;
  PetscFunctionReturn(0);
}

#undef __FUNCT__
#define __FUNCT__ "TaoLinearSolverPetsc::MinQuadraticTrustRegion"
int  TaoLinearSolverPetsc::MinQuadraticTrustRegion(TaoVec*tb,TaoVec*tx,double r, TaoTruth*flg){
  int info;

  PetscFunctionBegin;
  info=KSPSetType(ksp,KSPQCG); CHKERRQ(info);
  info = this->SetTrustRadius(r);CHKERRQ(info);
  info = this->Solve(tb,tx,flg);CHKERRQ(info);
  info = tx->Negate(); CHKERRQ(info);
  PetscFunctionReturn(0);
}

#undef __FUNCT__
#define __FUNCT__ "TaoLinearSolverPetsc::GetNumberIterations"
int TaoLinearSolverPetsc::GetNumberIterations(int * iters){
  PetscFunctionBegin;
  *iters=linear_its;
  PetscFunctionReturn(0);
}

#undef __FUNCT__
#define __FUNCT__ "TaoLinearSolverPetsc::SetTolerances"
int TaoLinearSolverPetsc::SetTolerances(double rtol, double atol, double dtol, int maxits){
  int info;

  PetscFunctionBegin;
  info = KSPSetTolerances(ksp,rtol,atol,dtol,maxits);CHKERRQ(info);
  PetscFunctionReturn(0);
}

#undef __FUNCT__
#define __FUNCT__ "TaoLinearSolverPetsc::View"
int TaoLinearSolverPetsc::View(){
  int info;
  PetscFunctionBegin;
  info=KSPView(this->ksp,this->pkspviewer);CHKERRQ(info);
  PetscFunctionReturn(0);
}

#undef __FUNCT__
#define __FUNCT__ "TaoLinearSolverPetsc::SetOptions"
int TaoLinearSolverPetsc::SetOptions(){
  int info;
  PetscFunctionBegin;
  info=KSPSetFromOptions(ksp); CHKERRQ(info);
  PetscFunctionReturn(0);
}

#undef __FUNCT__
#define __FUNCT__ "KSPDuplicate"
int KSPDuplicate(KSP ksp1, KSP* kspout){
  const char *prefix;
  PetscReal rtol,atol,dtol;
  int maxits,info;
  KSPType ktype;
  KSP ksp2;
  PCType ptype;
  PC  pc1,pc2;
  PetscTruth flg;
  MPI_Comm comm;
  PetscFunctionBegin;
  
  info = PetscObjectGetComm((PetscObject)ksp1,&comm); CHKERRQ(info);
  info=KSPCreate(comm,&ksp2); CHKERRQ(info);

  info=KSPGetOptionsPrefix( ksp1,&prefix); CHKERRQ(info);
  info=KSPSetOptionsPrefix( ksp2, prefix); CHKERRQ(info);

  info=KSPGetType(ksp1,&ktype); CHKERRQ(info);
  if (ktype){
    info=KSPSetType(ksp2, ktype); CHKERRQ(info);
  }
  info=KSPGetTolerances(ksp1,&rtol,&atol,&dtol,&maxits); CHKERRQ(info);
  info=KSPSetTolerances(ksp2, rtol, atol, dtol, PETSC_DEFAULT); CHKERRQ(info);

  info=KSPGetPC( ksp1,&pc1); CHKERRQ(info);
  info=KSPGetPC( ksp2,&pc2); CHKERRQ(info);

  info=KSPSetFromOptions(ksp2); CHKERRQ(info);

  info=PCGetType(pc1,&ptype); CHKERRQ(info);
  if (ptype){
    info=PetscStrncmp(ptype,PCMG,3,&flg); CHKERRQ(info);
    if (flg){
      info=PCSetType(pc2,PCBJACOBI); CHKERRQ(info);
    } else {
      info=PCSetType(pc2, ptype); CHKERRQ(info);
    }
  }

  *kspout=ksp2;
  PetscFunctionReturn(0);
}

#undef __FUNCT__
#define __FUNCT__ "TaoLinearSolverPetsc::Duplicate"
int TaoLinearSolverPetsc::Duplicate(TaoLinearSolver**M){
  int info;
  KSP ksp1=this->GetKSP();
  KSP ksp2;

  PetscFunctionBegin;
  info = KSPDuplicate(ksp1,&ksp2); CHKERRQ(info);
  *M = new TaoLinearSolverPetsc(ksp2);
  PetscFunctionReturn(0);
}

#endif
