# include <cstdlib>
# include <iostream>

using namespace std;
// 
//  By including PETSCKSP.H, we automatically also include all the information
//  defining lower level PETSC objects, in this case:
//    petsc.h - base PETSc routines
//    petscvec.h - vectors
//    petscsys.h - system routines
//    petscmat.h - matrices
//    petscis.h - index sets
//    petscksp.h - Krylov subspace methods
//    petscviewer.h - viewers
//    petscpc.h - preconditioners
//
# include "petscksp.h"
//
//  I guess the following pair of lines are intended to tell PETSc 
//  that the main program is called "main".
//
# undef __FUNCT__
# define __FUNCT__ "main"
//
//  Defining this string, and pointing to it in the PetscInitialize call,
//  simply allows a user to invoke the program with a "-help" option, which
//  will print out the corresponding help message.
//
static char help[] = "Solves a tridiagonal linear system with KSP.\n\n";

int main ( int argc, char **args );

//****************************************************************************80

int main ( int argc, char **args )

//****************************************************************************80
//
//  Purpose:
//
//    MAIN demonstrates the use of PETSc's KSP package to solve a system.
//
//  Discussion:
//
//    This example should NOT be run in parallel.
//
//    The matrix A is the [-1, 2, -1] tridiagonal matrix.
//
//  Modified:
//
//    02 November 2005
//
//  Parameters:
//
//    Mat A, the matrix.
//
//    Vec b, the right hand side.
//
//    PetscErrorCode ierr, the error code returned by every PETSc routine.
//
//    KSP ksp, records which linear solver is to be used.
//
//    PetscReal norm, the norm of the solution error.
//
//    PC pc, records which preconditioner is to be used.
//
//    Vec u, the exact solution.
//
//    Vec x, the approximate solution.
//
{
  Mat A;
  Vec b;
  PetscInt col[3];
  PetscInt i;
  PetscErrorCode ierr;
  PetscInt its;
  KSP ksp;
  PetscInt n = 10;
  PetscScalar neg_one = -1.0;
  PetscReal norm;
  PetscScalar one = 1.0;
  PC pc;
  PetscMPIInt size;
  Vec u;
  PetscScalar value[3];
  Vec x;

  PetscInitialize ( &argc, &args, (char *)0, help );
//
//  Determine the number of processors involved.
//  Complain if there are more than 1.
//
  MPI_Comm_size ( PETSC_COMM_WORLD, &size );

  if ( size != 1 ) 
  {
    SETERRQ(1,"This is a uniprocessor example only!");
  }

  PetscOptionsGetInt ( PETSC_NULL, "-n", &n, PETSC_NULL );
//
//  Compute the matrix and right-hand-side vector that define the linear system, Ax = b.
//
//  Create vector X "from scratch".
//
  VecCreate ( PETSC_COMM_WORLD, &x );
  PetscObjectSetName ( (PetscObject) x, "Solution" );
  VecSetSizes ( x, PETSC_DECIDE, n );
  VecSetFromOptions ( x );
//
//  Form B and U by duplicating X.
//
  VecDuplicate ( x, &b );
  VecDuplicate ( x, &u );
// 
//  Create the matrix.  When using MatCreate(), the matrix format can
//  be specified at runtime.
//
  MatCreate ( PETSC_COMM_WORLD, &A );
  MatSetSizes ( A, PETSC_DECIDE, PETSC_DECIDE, n, n );
  MatSetFromOptions ( A );
// 
//   Assemble the matrix.
//
  i = 0; 
  col[0] = 0; 
  col[1] = 1; 
  value[0] = 2.0; 
  value[1] = -1.0;

  MatSetValues ( A, 1, &i, 2, col, value, INSERT_VALUES );

  value[0] = -1.0; 
  value[1] = 2.0; 
  value[2] = -1.0;

  for ( i = 1; i < n - 1; i++ )
  {
    col[0] = i-1; 
    col[1] = i; 
    col[2] = i+1;

    MatSetValues ( A, 1, &i, 3, col, value, INSERT_VALUES );
  }

  i = n - 1; 
  col[0] = n - 2; 
  col[1] = n - 1;
  value[0] = -1.0; 
  value[1] = 2.0; 

  MatSetValues ( A, 1, &i, 2, col, value, INSERT_VALUES );


  MatAssemblyBegin ( A, MAT_FINAL_ASSEMBLY );
  MatAssemblyEnd ( A, MAT_FINAL_ASSEMBLY );
// 
//  Set the exact solution U = [1,1,...,1]..
//
  VecSet ( u, one );
// 
//  Compute right-hand-side vector b = A * U.
//
  MatMult ( A, u, b );
//
//  Create the linear solver and set various options
//
//  Create linear solver context
//
  KSPCreate ( PETSC_COMM_WORLD, &ksp );
// 
//  Set operators. Here the matrix that defines the linear system
//  also serves as the preconditioning matrix.
//
  KSPSetOperators ( ksp, A, A, DIFFERENT_NONZERO_PATTERN );
// 
//  Set linear solver defaults for this problem (optional).
//  - By extracting the KSP and PC contexts from the KSP context,
//    we can then directly call any KSP and PC routines to set
//    various options.
//
  KSPGetPC ( ksp, &pc );
  PCSetType ( pc, PCJACOBI );
  KSPSetTolerances ( ksp, 1.0e-7, PETSC_DEFAULT, PETSC_DEFAULT, PETSC_DEFAULT );
// 
//  Set runtime options, e.g.,
//    -ksp_type <type> -pc_type <type> -ksp_monitor -ksp_rtol <rtol>
//  These options will override those specified above as long as
//  KSPSetFromOptions() is called _after_ any other customization routines.
//
  KSPSetFromOptions ( ksp );
//
//  Solve the linear system.
//
  KSPSolve ( ksp, b, x );
// 
//  View solver info; we could instead use the option -ksp_view to
//  print this info to the screen at the conclusion of KSPSolve().
//
  KSPView ( ksp, PETSC_VIEWER_STDOUT_WORLD );
//
//  Check solution:
//  X <- X - U.
//  Norm <- || X ||
//  its <- number of iterations.
//
  VecAXPY ( x, neg_one, u );
  VecNorm ( x, NORM_2, &norm );
  KSPGetIterationNumber ( ksp, &its );
  PetscPrintf ( PETSC_COMM_WORLD, "Norm of error %A, Iterations %D\n",
    norm, its );
// 
//  Free work space.
//
  VecDestroy ( x );
  VecDestroy ( u );
  VecDestroy ( b ); 
  MatDestroy ( A );
  KSPDestroy ( ksp );
//
//  Close down PETSc.
//
  PetscFinalize ( );

  return 0;
}