!
!  Description: This example solves a nonlinear system on 1 processor with SNES.
!  We solve the  Bratu (SFI - solid fuel ignition) problem in a 2D rectangular
!  domain.  The command line options include:
!    -par <parameter>, where <parameter> indicates the nonlinearity of the problem
!       problem SFI:  <parameter> = Bratu parameter (0 <= par <= 6.81)
!    -mx <xg>, where <xg> = number of grid points in the x-direction
!    -my <yg>, where <yg> = number of grid points in the y-direction
!

!
!  --------------------------------------------------------------------------
!
!  Solid Fuel Ignition (SFI) problem.  This problem is modeled by
!  the partial differential equation
!
!          -Laplacian u - lambda*exp(u) = 0,  0 < x,y < 1,
!
!  with boundary conditions
!
!           u = 0  for  x = 0, x = 1, y = 0, y = 1.
!
!  A finite difference approximation with the usual 5-point stencil
!  is used to discretize the boundary value problem to obtain a nonlinear
!  system of equations.
!
!  The parallel version of this code is snes/tutorials/ex5f.F
!
!  --------------------------------------------------------------------------
#include <petsc/finclude/petscsnes.h>
#include <petsc/finclude/petscdraw.h>
module ex1fmodule
  use petscsnes
  implicit none
contains
  subroutine postcheck(snes, x, y, w, changed_y, changed_w, ctx, ierr)
    SNES snes
    PetscReal norm
    Vec tmp, x, y, w
    PetscBool changed_w, changed_y
    PetscErrorCode ierr
    PetscInt ctx
    PetscScalar mone
    MPIU_Comm comm

    character(len=PETSC_MAX_PATH_LEN) :: outputString

    PetscCallA(PetscObjectGetComm(snes, comm, ierr))
    PetscCallA(VecDuplicate(x, tmp, ierr))
    mone = -1.0
    PetscCallA(VecWAXPY(tmp, mone, x, w, ierr))
    PetscCallA(VecNorm(tmp, NORM_2, norm, ierr))
    PetscCallA(VecDestroy(tmp, ierr))
    write (outputString, *) norm
    PetscCallA(PetscPrintf(comm, 'Norm of search step '//trim(outputString)//'\n', ierr))
  end

! ---------------------------------------------------------------------
!
!  FormInitialGuess - Forms initial approximation.
!
!  Input Parameter:
!  X - vector
!
!  Output Parameters:
!  X - vector
!  ierr - error code
!
!  Notes:
!  This routine serves as a wrapper for the lower-level routine
!  "ApplicationInitialGuess", where the actual computations are
!  done using the standard Fortran style of treating the local
!  vector data as a multidimensional array over the local mesh.
!  This routine merely accesses the local vector data via
!  VecGetArray() and VecRestoreArray().
!
  subroutine FormInitialGuess(X, ierr)

!  Input/output variables:
    Vec X
    PetscErrorCode ierr

!     Declarations for use with local arrays:
    PetscScalar, pointer :: lx_v(:)

    ierr = 0

!  Get a pointer to vector data.
!    - VecGetArray() returns a pointer to the data array.
!    - You MUST call VecRestoreArray() when you no longer need access to
!      the array.

    PetscCallA(VecGetArray(X, lx_v, ierr))

!  Compute initial guess

    PetscCallA(ApplicationInitialGuess(lx_v, ierr))

!  Restore vector

    PetscCallA(VecRestoreArray(X, lx_v, ierr))

  end

!  ApplicationInitialGuess - Computes initial approximation, called by
!  the higher level routine FormInitialGuess().
!
!  Input Parameter:
!  x - local vector data
!
!  Output Parameters:
!  f - local vector data, f(x)
!  ierr - error code
!
!  Notes:
!  This routine uses standard Fortran-style computations over a 2-dim array.
!
  subroutine ApplicationInitialGuess(x, ierr)

!  Common blocks:
    PetscReal lambda
    PetscInt mx, my
    PetscBool fd_coloring
    common/params/lambda, mx, my, fd_coloring

!  Input/output variables:
    PetscScalar x(mx, my)
    PetscErrorCode ierr

!  Local variables:
    PetscInt i, j
    PetscReal temp1, temp, hx, hy, one

!  Set parameters

    ierr = 0
    one = 1.0
    hx = one/(mx - 1)
    hy = one/(my - 1)
    temp1 = lambda/(lambda + one)

    do j = 1, my
      temp = min(j - 1, my - j)*hy
      do i = 1, mx
        if (i == 1 .or. j == 1 .or. i == mx .or. j == my) then
          x(i, j) = 0.0
        else
          x(i, j) = temp1*sqrt(min(min(i - 1, mx - i)*hx, temp))
        end if
      end do
    end do

  end

! ---------------------------------------------------------------------
!
!  FormFunction - Evaluates nonlinear function, F(x).
!
!  Input Parameters:
!  snes  - the SNES context
!  X     - input vector
!  dummy - optional user-defined context, as set by SNESSetFunction()
!          (not used here)
!
!  Output Parameter:
!  F     - vector with newly computed function
!
!  Notes:
!  This routine serves as a wrapper for the lower-level routine
!  "ApplicationFunction", where the actual computations are
!  done using the standard Fortran style of treating the local
!  vector data as a multidimensional array over the local mesh.
!  This routine merely accesses the local vector data via
!  VecGetArray() and VecRestoreArray().
!
  subroutine FormFunction(snes, X, F, fdcoloring, ierr)

!  Input/output variables:
    SNES snes
    Vec X, F
    PetscErrorCode ierr
    MatFDColoring fdcoloring

!  Common blocks:
    PetscReal lambda
    PetscInt mx, my
    PetscBool fd_coloring
    common/params/lambda, mx, my, fd_coloring

!  Declarations for use with local arrays:
    PetscScalar, pointer :: lx_v(:), lf_v(:)
    PetscInt, pointer :: indices(:)

!  Get pointers to vector data.
!    - VecGetArray() returns a pointer to the data array.
!    - You MUST call VecRestoreArray() when you no longer need access to
!      the array.

    PetscCallA(VecGetArrayRead(X, lx_v, ierr))
    PetscCallA(VecGetArray(F, lf_v, ierr))

!  Compute function

    PetscCallA(ApplicationFunction(lx_v, lf_v, ierr))

!  Restore vectors

    PetscCallA(VecRestoreArrayRead(X, lx_v, ierr))
    PetscCallA(VecRestoreArray(F, lf_v, ierr))

    PetscCallA(PetscLogFlops(11.0d0*mx*my, ierr))
!
!     fdcoloring is in the common block and used here ONLY to test the
!     calls to MatFDColoringGetPerturbedColumns() and  MatFDColoringRestorePerturbedColumns()
!
    if (fd_coloring) then
      PetscCallA(MatFDColoringGetPerturbedColumns(fdcoloring, PETSC_NULL_INTEGER, indices, ierr))
      print *, 'Indices from GetPerturbedColumns'
      write (*, 1000) indices
1000  format(50i4)
      PetscCallA(MatFDColoringRestorePerturbedColumns(fdcoloring, PETSC_NULL_INTEGER, indices, ierr))
    end if
  end

! ---------------------------------------------------------------------
!
!  ApplicationFunction - Computes nonlinear function, called by
!  the higher level routine FormFunction().
!
!  Input Parameter:
!  x    - local vector data
!
!  Output Parameters:
!  f    - local vector data, f(x)
!  ierr - error code
!
!  Notes:
!  This routine uses standard Fortran-style computations over a 2-dim array.
!
  subroutine ApplicationFunction(x, f, ierr)

!  Common blocks:
    PetscReal lambda
    PetscInt mx, my
    PetscBool fd_coloring
    common/params/lambda, mx, my, fd_coloring

!  Input/output variables:
    PetscScalar x(mx, my), f(mx, my)
    PetscErrorCode ierr

!  Local variables:
    PetscScalar two, one, hx, hy
    PetscScalar hxdhy, hydhx, sc
    PetscScalar u, uxx, uyy
    PetscInt i, j

    ierr = 0
    one = 1.0
    two = 2.0
    hx = one/(mx - 1)
    hy = one/(my - 1)
    sc = hx*hy*lambda
    hxdhy = hx/hy
    hydhx = hy/hx

!  Compute function

    do j = 1, my
      do i = 1, mx
        if (i == 1 .or. j == 1 .or. i == mx .or. j == my) then
          f(i, j) = x(i, j)
        else
          u = x(i, j)
          uxx = hydhx*(two*u - x(i - 1, j) - x(i + 1, j))
          uyy = hxdhy*(two*u - x(i, j - 1) - x(i, j + 1))
          f(i, j) = uxx + uyy - sc*exp(u)
        end if
      end do
    end do

  end

! ---------------------------------------------------------------------
!
!  FormJacobian - Evaluates Jacobian matrix.
!
!  Input Parameters:
!  snes    - the SNES context
!  x       - input vector
!  dummy   - optional user-defined context, as set by SNESSetJacobian()
!            (not used here)
!
!  Output Parameters:
!  jac      - Jacobian matrix
!  jac_prec - optionally different matrix used to construct the preconditioner (not used here)
!
!  Notes:
!  This routine serves as a wrapper for the lower-level routine
!  "ApplicationJacobian", where the actual computations are
!  done using the standard Fortran style of treating the local
!  vector data as a multidimensional array over the local mesh.
!  This routine merely accesses the local vector data via
!  VecGetArray() and VecRestoreArray().
!
  subroutine FormJacobian(snes, X, jac, jac_prec, dummy, ierr)

!  Input/output variables:
    SNES snes
    Vec X
    Mat jac, jac_prec
    PetscErrorCode ierr
    integer dummy

!  Common blocks:
    PetscReal lambda
    PetscInt mx, my
    PetscBool fd_coloring
    common/params/lambda, mx, my, fd_coloring

!  Declarations for use with local array:
    PetscScalar, pointer :: lx_v(:)

!  Get a pointer to vector data

    PetscCallA(VecGetArrayRead(X, lx_v, ierr))

!  Compute Jacobian entries

    PetscCallA(ApplicationJacobian(lx_v, jac, jac_prec, ierr))

!  Restore vector

    PetscCallA(VecRestoreArrayRead(X, lx_v, ierr))

!  Assemble matrix

    PetscCallA(MatAssemblyBegin(jac_prec, MAT_FINAL_ASSEMBLY, ierr))
    PetscCallA(MatAssemblyEnd(jac_prec, MAT_FINAL_ASSEMBLY, ierr))

  end

! ---------------------------------------------------------------------
!
!  ApplicationJacobian - Computes Jacobian matrix, called by
!  the higher level routine FormJacobian().
!
!  Input Parameters:
!  x        - local vector data
!
!  Output Parameters:
!  jac      - Jacobian matrix
!  jac_prec - optionally different matrix used to construct the preconditioner (not used here)
!  ierr     - error code
!
!  Notes:
!  This routine uses standard Fortran-style computations over a 2-dim array.
!
  subroutine ApplicationJacobian(x, jac, jac_prec, ierr)
!  Common blocks:
    PetscReal lambda
    PetscInt mx, my
    PetscBool fd_coloring
    common/params/lambda, mx, my, fd_coloring

!  Input/output variables:
    PetscScalar x(mx, my)
    Mat jac, jac_prec
    PetscErrorCode ierr

!  Local variables:
    PetscInt i, j, row(1), col(5), i1, i5
    PetscScalar two, one, hx, hy
    PetscScalar hxdhy, hydhx, sc, v(5)

!  Set parameters

    i1 = 1
    i5 = 5
    one = 1.0
    two = 2.0
    hx = one/(mx - 1)
    hy = one/(my - 1)
    sc = hx*hy
    hxdhy = hx/hy
    hydhx = hy/hx

!  Compute entries of the Jacobian matrix
!   - Here, we set all entries for a particular row at once.
!   - Note that MatSetValues() uses 0-based row and column numbers
!     in Fortran as well as in C.

    do j = 1, my
      row(1) = (j - 1)*mx - 1
      do i = 1, mx
        row(1) = row(1) + 1
!           boundary points
        if (i == 1 .or. j == 1 .or. i == mx .or. j == my) then
          PetscCallA(MatSetValues(jac_prec, i1, row, i1, row, [one], INSERT_VALUES, ierr))
!           interior grid points
        else
          v(1) = -hxdhy
          v(2) = -hydhx
          v(3) = two*(hydhx + hxdhy) - sc*lambda*exp(x(i, j))
          v(4) = -hydhx
          v(5) = -hxdhy
          col(1) = row(1) - mx
          col(2) = row(1) - 1
          col(3) = row(1)
          col(4) = row(1) + 1
          col(5) = row(1) + mx
          PetscCallA(MatSetValues(jac_prec, i1, row, i5, col, v, INSERT_VALUES, ierr))
        end if
      end do
    end do

  end

end module ex1fmodule
program main
  use petscdraw
  use petscsnes
  use ex1fmodule
  implicit none
!
! - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
!                   Variable declarations
! - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
!
!  Variables:
!     snes        - nonlinear solver
!     x, r        - solution, residual vectors
!     J           - Jacobian matrix
!     its         - iterations for convergence
!     matrix_free - flag - 1 indicates matrix-free version
!     lambda      - nonlinearity parameter
!     draw        - drawing context
!
  SNES snes
  MatColoring mc
  Vec x, r
  PetscDraw draw
  Mat J
  PetscBool matrix_free, flg, fd_coloring
  PetscErrorCode ierr
  PetscInt its, N, mx, my, i5
  PetscMPIInt size, rank
  PetscReal lambda_max, lambda_min, lambda
  MatFDColoring fdcoloring
  ISColoring iscoloring
  PetscBool pc
  integer4 imx, imy
  character(len=PETSC_MAX_PATH_LEN) :: outputString
  PetscScalar, pointer :: lx_v(:)
  integer4 xl, yl, width, height

!  Store parameters in common block

  common/params/lambda, mx, my, fd_coloring

! - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
!  Initialize program
! - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -

  PetscCallA(PetscInitialize(ierr))
  PetscCallMPIA(MPI_Comm_size(PETSC_COMM_WORLD, size, ierr))
  PetscCallMPIA(MPI_Comm_rank(PETSC_COMM_WORLD, rank, ierr))

  PetscCheckA(size == 1, PETSC_COMM_SELF, PETSC_ERR_WRONG_MPI_SIZE, 'This is a uniprocessor example only')

!  Initialize problem parameters
  i5 = 5
  lambda_max = 6.81
  lambda_min = 0.0
  lambda = 6.0
  mx = 4
  my = 4
  PetscCallA(PetscOptionsGetInt(PETSC_NULL_OPTIONS, PETSC_NULL_CHARACTER, '-mx', mx, flg, ierr))
  PetscCallA(PetscOptionsGetInt(PETSC_NULL_OPTIONS, PETSC_NULL_CHARACTER, '-my', my, flg, ierr))
  PetscCallA(PetscOptionsGetReal(PETSC_NULL_OPTIONS, PETSC_NULL_CHARACTER, '-par', lambda, flg, ierr))
  PetscCheckA(lambda < lambda_max .and. lambda > lambda_min, PETSC_COMM_SELF, PETSC_ERR_USER, 'Lambda out of range ')
  N = mx*my
  pc = PETSC_FALSE
  PetscCallA(PetscOptionsGetBool(PETSC_NULL_OPTIONS, PETSC_NULL_CHARACTER, '-pc', pc, PETSC_NULL_BOOL, ierr))

! - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
!  Create nonlinear solver context
! - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -

  PetscCallA(SNESCreate(PETSC_COMM_WORLD, snes, ierr))

  if (pc .eqv. PETSC_TRUE) then
    PetscCallA(SNESSetType(snes, SNESNEWTONTR, ierr))
    PetscCallA(SNESNewtonTRSetPostCheck(snes, postcheck, snes, ierr))
  end if

! - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
!  Create vector data structures; set function evaluation routine
! - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -

  PetscCallA(VecCreate(PETSC_COMM_WORLD, x, ierr))
  PetscCallA(VecSetSizes(x, PETSC_DECIDE, N, ierr))
  PetscCallA(VecSetFromOptions(x, ierr))
  PetscCallA(VecDuplicate(x, r, ierr))

!  Set function evaluation routine and vector.  Whenever the nonlinear
!  solver needs to evaluate the nonlinear function, it will call this
!  routine.
!   - Note that the final routine argument is the user-defined
!     context that provides application-specific data for the
!     function evaluation routine.

  PetscCallA(SNESSetFunction(snes, r, FormFunction, fdcoloring, ierr))

! - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
!  Create matrix data structure; set Jacobian evaluation routine
! - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -

!  Create matrix. Here we only approximately preallocate storage space
!  for the Jacobian.  See the users manual for a discussion of better
!  techniques for preallocating matrix memory.

  PetscCallA(PetscOptionsHasName(PETSC_NULL_OPTIONS, PETSC_NULL_CHARACTER, '-snes_mf', matrix_free, ierr))
  if (.not. matrix_free) then
    PetscCallA(MatCreateSeqAIJ(PETSC_COMM_WORLD, N, N, i5, PETSC_NULL_INTEGER_ARRAY, J, ierr))
  end if

!
!     This option will cause the Jacobian to be computed via finite differences
!    efficiently using a coloring of the columns of the matrix.
!
  fd_coloring = .false.
  PetscCallA(PetscOptionsHasName(PETSC_NULL_OPTIONS, PETSC_NULL_CHARACTER, '-snes_fd_coloring', fd_coloring, ierr))
  if (fd_coloring) then

!
!      This initializes the nonzero structure of the Jacobian. This is artificial
!      because clearly if we had a routine to compute the Jacobian we won't need
!      to use finite differences.
!
    PetscCallA(FormJacobian(snes, x, J, J, 0, ierr))
!
!       Color the matrix, i.e. determine groups of columns that share no common
!      rows. These columns in the Jacobian can all be computed simultaneously.
!
    PetscCallA(MatColoringCreate(J, mc, ierr))
    PetscCallA(MatColoringSetType(mc, MATCOLORINGNATURAL, ierr))
    PetscCallA(MatColoringSetFromOptions(mc, ierr))
    PetscCallA(MatColoringApply(mc, iscoloring, ierr))
    PetscCallA(MatColoringDestroy(mc, ierr))
!
!       Create the data structure that SNESComputeJacobianDefaultColor() uses
!       to compute the actual Jacobians via finite differences.
!
    PetscCallA(MatFDColoringCreate(J, iscoloring, fdcoloring, ierr))
    PetscCallA(MatFDColoringSetFunction(fdcoloring, FormFunction, fdcoloring, ierr))
    PetscCallA(MatFDColoringSetFromOptions(fdcoloring, ierr))
    PetscCallA(MatFDColoringSetUp(J, iscoloring, fdcoloring, ierr))
!
!        Tell SNES to use the routine SNESComputeJacobianDefaultColor()
!      to compute Jacobians.
!
    PetscCallA(SNESSetJacobian(snes, J, J, SNESComputeJacobianDefaultColor, fdcoloring, ierr))
    PetscCallA(ISColoringDestroy(iscoloring, ierr))

  else if (.not. matrix_free) then

!  Set Jacobian matrix data structure and default Jacobian evaluation
!  routine.  Whenever the nonlinear solver needs to compute the
!  Jacobian matrix, it will call this routine.
!   - Note that the final routine argument is the user-defined
!     context that provides application-specific data for the
!     Jacobian evaluation routine.
!   - The user can override with:
!      -snes_fd : default finite differencing approximation of Jacobian
!      -snes_mf : matrix-free Newton-Krylov method with no preconditioning
!                 (unless user explicitly sets preconditioner)
!      -snes_mf_operator : form matrix from which to construct the preconditioner as set by the user,
!                          but use matrix-free approx for Jacobian-vector
!                          products within Newton-Krylov method
!
    PetscCallA(SNESSetJacobian(snes, J, J, FormJacobian, 0, ierr))
  end if

! - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
!  Customize nonlinear solver; set runtime options
! - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -

!  Set runtime options (e.g., -snes_monitor -snes_rtol <rtol> -ksp_type <type>)

  PetscCallA(SNESSetFromOptions(snes, ierr))

! - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
!  Evaluate initial guess; then solve nonlinear system.
! - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -

!  Note: The user should initialize the vector, x, with the initial guess
!  for the nonlinear solver prior to calling SNESSolve().  In particular,
!  to employ an initial guess of zero, the user should explicitly set
!  this vector to zero by calling VecSet().

  PetscCallA(FormInitialGuess(x, ierr))
  PetscCallA(SNESSolve(snes, PETSC_NULL_VEC, x, ierr))
  PetscCallA(SNESGetIterationNumber(snes, its, ierr))
  write (outputString, *) its
  PetscCallA(PetscPrintf(PETSC_COMM_WORLD, 'Number of SNES iterations = '//trim(outputString)//'\n', ierr))

!  PetscDraw contour plot of solution

  xl = 300
  yl = 0
  width = 300
  height = 300
  PetscCallA(PetscDrawCreate(PETSC_COMM_WORLD, PETSC_NULL_CHARACTER, 'Solution', xl, yl, width, height, draw, ierr))
  PetscCallA(PetscDrawSetFromOptions(draw, ierr))

  PetscCallA(VecGetArrayRead(x, lx_v, ierr))
  imx = int(mx, kind=kind(imx))
  imy = int(my, kind=kind(imy))
  PetscCallA(PetscDrawTensorContour(draw, imx, imy, PETSC_NULL_REAL_ARRAY, PETSC_NULL_REAL_ARRAY, lx_v, ierr))
  PetscCallA(VecRestoreArrayRead(x, lx_v, ierr))

! - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
!  Free work space.  All PETSc objects should be destroyed when they
!  are no longer needed.
! - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -

  if (.not. matrix_free) PetscCallA(MatDestroy(J, ierr))
  if (fd_coloring) PetscCallA(MatFDColoringDestroy(fdcoloring, ierr))

  PetscCallA(VecDestroy(x, ierr))
  PetscCallA(VecDestroy(r, ierr))
  PetscCallA(SNESDestroy(snes, ierr))
  PetscCallA(PetscDrawDestroy(draw, ierr))
  PetscCallA(PetscFinalize(ierr))
end
!
!/*TEST
!
!   build:
!      requires: !single !complex
!
!   test:
!      args: -snes_monitor_short -nox -snes_type newtontr -ksp_gmres_cgs_refinement_type refine_always
!
!   test:
!      suffix: 2
!      args: -snes_monitor_short -nox -snes_fd -ksp_gmres_cgs_refinement_type refine_always
!
!   test:
!      suffix: 3
!      args: -snes_monitor_short -nox -snes_fd_coloring -mat_coloring_type sl -ksp_gmres_cgs_refinement_type refine_always
!      filter: sort -b
!      filter_output: sort -b
!
!   test:
!     suffix: 4
!     args: -pc -par 6.807 -nox
!
!TEST*/