!
!  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
!
!  --------------------------------------------------------------------------
      subroutine postcheck(snes, x, y, w, changed_y, changed_w, ctx, ierr)
#include <petsc/finclude/petscsnes.h>
        use petscsnes
        implicit none
        SNES snes
        PetscReal norm
        Vec tmp, x, y, w
        PetscBool changed_w, changed_y
        PetscErrorCode ierr
        PetscInt ctx
        PetscScalar mone
        MPI_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

      program main
#include <petsc/finclude/petscdraw.h>
        use petscdraw
        use petscsnes
        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
        external postcheck
        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

!  Note: Any user-defined Fortran routines (such as FormJacobian)
!  MUST be declared as external.

        external FormFunction, FormInitialGuess, FormJacobian

! - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
!  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(SNESGetJacobian(snes, J, PETSC_NULL_MAT, PETSC_NULL_FUNCTION, PETSC_NULL_INTEGER, ierr))
          PetscCallA(SNESGetJacobian(snes, J, PETSC_NULL_MAT, PETSC_NULL_FUNCTION, 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

! ---------------------------------------------------------------------
!
!  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)
        use petscsnes
        implicit none

!  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)
        use petscksp
        implicit none

!  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 20 j = 1, my
          temp = min(j - 1, my - j)*hy
          do 10 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
10          continue
20          continue

          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)
            use petscsnes
            implicit none

!  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)
            use petscsnes
            implicit none

!  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 20 j = 1, my
              do 10 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
10              continue
20              continue

              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)
                use petscsnes
                implicit none

!  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)
                use petscsnes
                implicit none

!  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 20 j = 1, my
                  row(1) = (j - 1)*mx - 1
                  do 10 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
10                  continue
20                  continue

                  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*/