!-----------------------------------------------------------------
! Programmer(s): Daniel R. Reynolds @ SMU
!-----------------------------------------------------------------
! LLNS/SMU Copyright Start
! Copyright (c) 2015, Southern Methodist University and 
! Lawrence Livermore National Security
!
! This work was performed under the auspices of the U.S. Department 
! of Energy by Southern Methodist University and Lawrence Livermore 
! National Laboratory under Contract DE-AC52-07NA27344.
! Produced at Southern Methodist University and the Lawrence 
! Livermore National Laboratory.
!
! All rights reserved.
! For details, see the LICENSE file.
! LLNS/SMU Copyright End
! Copyright (c) 2013, Southern Methodist University.
! All rights reserved.
! For details, see the LICENSE file.
!-----------------------------------------------------------------
! Example problem:
! 
! The following test simulates a simple anisotropic 2D heat 
! equation,
!    u_t = kx*u_xx + ky*u_yy + h,
! for t in [0, 10], (x,y) in [0, 1]^2, with initial conditions
!    u(0,x,y) =  0,
! stationary boundary conditions, i.e. 
!    u_t(t,0,y) = u_t(t,1,y) = u_t(t,x,0) = u_t(t,x,1) = 0,
! and a heat source of the form
!    h(x,y) = sin(pi*x)*sin(2*pi*y).
!
! Under this setup, the problem has an analytical solution:
!    u(t,x,y) = a(t)*sin(pi*x)*sin(2*pi*y), where
!    a(t) = (1 - exp(-(kx+4*ky)*pi^2*t)) / ((kx+4*ky)*pi^2).
! 
! The spatial derivatives are computed using second-order 
! centered differences, with the data distributed over nx*ny
! points on a uniform spatial grid.
!
! The spatial grid parameters nx and ny, the parameters kx and ky, 
! as well as the desired relative and absolute solver tolerances, 
! are provided in the input file input_heat2D.txt.
! 
! This program solves the problem with a DIRK method.  This 
! employs a Newton iteration with the PCG iterative linear solver, 
! which itself uses a Jacobi preconditioner.  The example uses the 
! built-in finite-difference Jacobian-vector product routine, but 
! supplies both the RHS and preconditioner setup/solve functions.
!
! 20 outputs are printed at equal intervals, and run statistics 
! are printed at the end.
!-----------------------------------------------------------------

!-----------------------------------------------------------------
! Main driver program
!-----------------------------------------------------------------
module UserData
  !---------------------------------------------------------------
  ! Description: 
  !    Module containing problem-defining parameters, as well as
  !    data buffers for MPI exchanges with neighboring processes.
  !    Also contains routines to:
  !      (a) initialize the module
  !      (b) perform exchanges
  !      (c) free module data.
  !---------------------------------------------------------------
  implicit none
  save

  integer*8 :: nx          ! global number of x grid points 
  integer*8 :: ny          ! global number of y grid points 
  integer*8 :: is          ! global x indices of this subdomain
  integer*8 :: ie
  integer*8 :: js          ! global y indices of this subdomain
  integer*8 :: je
  integer*8 :: nxl         ! local number of x grid points 
  integer*8 :: nyl         ! local number of y grid points 
  real*8    :: dx          ! x-directional mesh spacing 
  real*8    :: dy          ! y-directional mesh spacing 
  real*8    :: kx          ! x-directional diffusion coefficient 
  real*8    :: ky          ! y-directional diffusion coefficient 
  real*8, dimension(:,:), allocatable :: h    ! heat source vector
  real*8, dimension(:,:), allocatable :: d    ! inverse of Jacobian diagonal
  integer :: comm                             ! communicator object
  integer :: myid                             ! MPI process ID
  integer :: nprocs                           ! total number of MPI processes
  logical :: HaveBdry(2,2)                    ! flags denoting if on physical boundary
  real*8, dimension(:), allocatable :: Erecv  ! receive buffers for neighbor exchange
  real*8, dimension(:), allocatable :: Wrecv
  real*8, dimension(:), allocatable :: Nrecv
  real*8, dimension(:), allocatable :: Srecv
  real*8, dimension(:), allocatable :: Esend  ! send buffers for neighbor exchange
  real*8, dimension(:), allocatable :: Wsend
  real*8, dimension(:), allocatable :: Nsend
  real*8, dimension(:), allocatable :: Ssend

contains

  !---------------------------------------------------------------
  ! Initialize memory allocated within Userdata (set to defaults)
  !---------------------------------------------------------------
  subroutine InitUserData()
    implicit none
    include "mpif.h"
    nx = 0
    ny = 0
    is = 0
    ie = 0
    js = 0
    je = 0
    nxl = 0
    nyl = 0
    dx = 0.d0
    dy = 0.d0
    kx = 0.d0
    ky = 0.d0
    if (allocated(h))  deallocate(h)
    if (allocated(d))  deallocate(d)
    comm = MPI_COMM_WORLD
    myid = 0
    nprocs = 0
    HaveBdry = .true.
    if (allocated(Erecv))  deallocate(Erecv)
    if (allocated(Wrecv))  deallocate(Wrecv)
    if (allocated(Nrecv))  deallocate(Nrecv)
    if (allocated(Srecv))  deallocate(Srecv)
    if (allocated(Esend))  deallocate(Esend)
    if (allocated(Wsend))  deallocate(Wsend)
    if (allocated(Nsend))  deallocate(Nsend)
    if (allocated(Ssend))  deallocate(Ssend)
  end subroutine InitUserData
  !---------------------------------------------------------------


  !---------------------------------------------------------------
  ! Set up parallel decomposition
  !---------------------------------------------------------------
  subroutine SetupDecomp(ierr)
    ! declarations
    implicit none
    include "mpif.h"
    integer, intent(out) :: ierr
    integer :: dims(2), periods(2), coords(2)
    
    ! internals

    ! check that this has not been called before
    if (allocated(h) .or. allocated(d)) then
       write(0,*) "SetupDecomp warning: parallel decomposition already set up"
       ierr = 1
       return
    end if

    ! get suggested parallel decomposition
    dims = (/0, 0/)
    call MPI_Comm_size(MPI_COMM_WORLD, nprocs, ierr)
    if (ierr /= MPI_SUCCESS) then
       write(0,*) "Error in MPI_Comm_size = " , ierr
       return
    end if
    call MPI_Dims_create(nprocs, 2, dims, ierr)
    if (ierr /= MPI_SUCCESS) then
       write(0,*) "Error in MPI_Dims_create = " , ierr
       return
    end if

    ! set up 2D Cartesian communicator
    periods = (/0, 0/)
    call MPI_Cart_create(MPI_COMM_WORLD, 2, dims, periods, 0, comm, ierr)
    if (ierr /= MPI_SUCCESS) then
       write(0,*) "Error in MPI_Cart_create = " , ierr
       return
    end if
    call MPI_Comm_rank(comm, myid, ierr)
    if (ierr /= MPI_SUCCESS) then
       write(0,*) "Error in MPI_Comm_rank = " , ierr
       return
    end if

    ! determine local extents
    call MPI_Cart_get(comm, 2, dims, periods, coords, ierr)
    if (ierr /= MPI_SUCCESS) then
       write(0,*) "Error in MPI_Cart_get = " , ierr
       return
    end if
    is = nx*coords(1)/dims(1) + 1
    ie = nx*(coords(1)+1)/dims(1)
    js = ny*coords(2)/dims(2) + 1
    je = ny*(coords(2)+1)/dims(2)
    nxl = ie-is+1
    nyl = je-js+1

    ! determine if I have neighbors, and allocate exchange buffers
    HaveBdry(1,1) = (is == 1)
    HaveBdry(1,2) = (ie == nx)
    HaveBdry(2,1) = (js == 1)
    HaveBdry(2,2) = (je == ny)
    if (.not. HaveBdry(1,1)) then
       allocate(Wrecv(nyl))
       allocate(Wsend(nyl))
    endif
    if (.not. HaveBdry(1,2)) then
       allocate(Erecv(nyl))
       allocate(Esend(nyl))
    endif
    if (.not. HaveBdry(2,1)) then
       allocate(Srecv(nxl))
       allocate(Ssend(nxl))
    endif
    if (.not. HaveBdry(2,2)) then
       allocate(Nrecv(nxl))
       allocate(Nsend(nxl))
    endif

    ! allocate temporary vectors
    allocate(h(nxl,nyl))    ! Create vector for heat source
    allocate(d(nxl,nyl))    ! Create vector for Jacobian diagonal
    
    ierr = 0     ! return with success flag
    return
  end subroutine SetupDecomp
  !---------------------------------------------------------------


  !---------------------------------------------------------------
  ! Perform neighbor exchange
  !---------------------------------------------------------------
  subroutine Exchange(y, ierr)
    ! declarations
    implicit none
    include "mpif.h"
    real*8,  intent(in)  :: y(nxl,nyl)
    integer, intent(out) :: ierr
    integer :: reqSW, reqSE, reqSS, reqSN, reqRW, reqRE, reqRS, reqRN;
    integer :: stat(MPI_STATUS_SIZE)
    integer*8 :: i
    integer :: ipW, ipE, ipS, ipN
    integer :: coords(2), dims(2), periods(2), nbcoords(2)
    
    ! internals

    ! MPI neighborhood information
    call MPI_Cart_get(comm, 2, dims, periods, coords, ierr)
    if (ierr /= MPI_SUCCESS) then
       write(0,*) "Error in MPI_Cart_get = ", ierr
       return
    endif
    if (.not. HaveBdry(1,1)) then
       nbcoords = (/ coords(1)-1, coords(2) /)
       call MPI_Cart_rank(comm, nbcoords, ipW, ierr)
       if (ierr /= MPI_SUCCESS) then 
          write(0,*) "Error in MPI_Cart_rank = ", ierr
          return
       endif
    endif
    if (.not. HaveBdry(1,2)) then
       nbcoords = (/ coords(1)+1, coords(2) /)
       call MPI_Cart_rank(comm, nbcoords, ipE, ierr)
       if (ierr /= MPI_SUCCESS) then 
          write(0,*) "Error in MPI_Cart_rank = ", ierr
          return
       endif
    endif
    if (.not. HaveBdry(2,1)) then
       nbcoords = (/ coords(1), coords(2)-1 /)
       call MPI_Cart_rank(comm, nbcoords, ipS, ierr)
       if (ierr /= MPI_SUCCESS) then 
          write(0,*) "Error in MPI_Cart_rank = ", ierr
          return
       endif
    endif
    if (.not. HaveBdry(2,2)) then
       nbcoords = (/ coords(1), coords(2)+1 /)
       call MPI_Cart_rank(comm, nbcoords, ipN, ierr)
       if (ierr /= MPI_SUCCESS) then 
          write(0,*) "Error in MPI_Cart_rank = ", ierr
          return
       endif
    endif

    ! open Irecv buffers
    if (.not. HaveBdry(1,1)) then
       call MPI_Irecv(Wrecv, nyl, MPI_DOUBLE_PRECISION, ipW, &
                      MPI_ANY_TAG, comm, reqRW, ierr)
       if (ierr /= MPI_SUCCESS) then 
          write(0,*) "Error in MPI_Irecv = ", ierr
          return
       endif
    endif
    if (.not. HaveBdry(1,2)) then
       call MPI_Irecv(Erecv, nyl, MPI_DOUBLE_PRECISION, ipE, &
                      MPI_ANY_TAG, comm, reqRE, ierr)
       if (ierr /= MPI_SUCCESS) then 
          write(0,*) "Error in MPI_Irecv = ", ierr
          return
       endif
    endif
    if (.not. HaveBdry(2,1)) then
       call MPI_Irecv(Srecv, nxl, MPI_DOUBLE_PRECISION, ipS, &
                      MPI_ANY_TAG, comm, reqRS, ierr)
       if (ierr /= MPI_SUCCESS) then 
          write(0,*) "Error in MPI_Irecv = ", ierr
          return
       endif
    endif
    if (.not. HaveBdry(2,2)) then
       call MPI_Irecv(Nrecv, nxl, MPI_DOUBLE_PRECISION, ipN, &
                      MPI_ANY_TAG, comm, reqRN, ierr)
       if (ierr /= MPI_SUCCESS) then 
          write(0,*) "Error in MPI_Irecv = ", ierr
          return
       endif
    endif

    ! send data
    if (.not. HaveBdry(1,1)) then
       do i=1,nyl
          Wsend(i) = y(1,i)
       enddo
       call MPI_Isend(Wsend, nyl, MPI_DOUBLE_PRECISION, ipW, 0, &
                      comm, reqSW, ierr)
       if (ierr /= MPI_SUCCESS) then 
          write(0,*) "Error in MPI_Isend = ", ierr
          return
       endif
    endif
    if (.not. HaveBdry(1,2)) then
       do i=1,nyl
          Esend(i) = y(nxl,i)
       enddo
       call MPI_Isend(Esend, nyl, MPI_DOUBLE_PRECISION, ipE, 1, &
                      comm, reqSE, ierr)
       if (ierr /= MPI_SUCCESS) then 
          write(0,*) "Error in MPI_Isend = ", ierr
          return
       endif
    endif
    if (.not. HaveBdry(2,1)) then
       do i=0,nxl
          Ssend(i) = y(i,1)
       enddo
       call MPI_Isend(Ssend, nxl, MPI_DOUBLE_PRECISION, ipS, 2, &
                      comm, reqSS, ierr)
       if (ierr /= MPI_SUCCESS) then 
          write(0,*) "Error in MPI_Isend = ", ierr
          return
       endif
    endif
    if (.not. HaveBdry(2,2)) then
       do i=0,nxl
          Nsend(i) = y(i,nyl)
       enddo
       call MPI_Isend(Nsend, nxl, MPI_DOUBLE_PRECISION, ipN, 3, &
                      comm, reqSN, ierr)
       if (ierr /= MPI_SUCCESS) then 
          write(0,*) "Error in MPI_Isend = ", ierr
          return
       endif
    endif

    ! wait for messages to finish
    if (.not. HaveBdry(1,1)) then
       call MPI_Wait(reqRW, stat, ierr)
       if (ierr /= MPI_SUCCESS) then 
          write(0,*) "Error in MPI_Wait = ", ierr
          return
       endif
       call MPI_Wait(reqSW, stat, ierr)
       if (ierr /= MPI_SUCCESS) then 
          write(0,*) "Error in MPI_Wait = ", ierr
          return
       endif
    endif
    if (.not. HaveBdry(1,2)) then
       call MPI_Wait(reqRE, stat, ierr)
       if (ierr /= MPI_SUCCESS) then 
          write(0,*) "Error in MPI_Wait = ", ierr
          return
       endif
       call MPI_Wait(reqSE, stat, ierr)
       if (ierr /= MPI_SUCCESS) then 
          write(0,*) "Error in MPI_Wait = ", ierr
          return
       endif
    endif
    if (.not. HaveBdry(2,1)) then
       call MPI_Wait(reqRS, stat, ierr)
       if (ierr /= MPI_SUCCESS) then 
          write(0,*) "Error in MPI_Wait = ", ierr
          return
       endif
       call MPI_Wait(reqSS, stat, ierr)
       if (ierr /= MPI_SUCCESS) then 
          write(0,*) "Error in MPI_Wait = ", ierr
          return
       endif
    endif
    if (.not. HaveBdry(2,2)) then
       call MPI_Wait(reqRN, stat, ierr)
       if (ierr /= MPI_SUCCESS) then 
          write(0,*) "Error in MPI_Wait = ", ierr
          return
       endif
       call MPI_Wait(reqSN, stat, ierr)
       if (ierr /= MPI_SUCCESS) then 
          write(0,*) "Error in MPI_Wait = ", ierr
          return
       endif
    endif

    ierr = 0     ! return with success flag
    return
  end subroutine Exchange
  !---------------------------------------------------------------


  !---------------------------------------------------------------
  ! Free memory allocated within Userdata
  !---------------------------------------------------------------
  subroutine FreeUserData(ierr)
    implicit none
    integer, intent(out) :: ierr
    if (.not. HaveBdry(1,1)) then
       deallocate(Wrecv)
       deallocate(Wsend)
    endif
    if (.not. HaveBdry(1,2)) then
       deallocate(Erecv)
       deallocate(Esend)
    endif
    if (.not. HaveBdry(2,1)) then
       deallocate(Srecv)
       deallocate(Ssend)
    endif
    if (.not. HaveBdry(2,2)) then
       deallocate(Nrecv)
       deallocate(Nsend)
    endif
    ierr = 0     ! return with success flag
    return
  end subroutine FreeUserData
  !---------------------------------------------------------------


  !---------------------------------------------------------------
  ! RMS norm function for parallel array data
  !---------------------------------------------------------------
  subroutine PRMS(y,yrms,ierr)
    ! declarations
    implicit none
    include "mpif.h"
    integer, intent(out) :: ierr
    real*8,  intent(in)  :: y(nxl,nyl)
    real*8,  intent(out) :: yrms
    real*8 :: lsum, gsum

    ! internals
    lsum = sum(y**2)
    call MPI_Allreduce(lsum, gsum, 1, MPI_DOUBLE_PRECISION, &
                       MPI_SUM, comm, ierr)
    if (ierr /= MPI_SUCCESS) then
       write(0,*) "Error in MPI_Allreduce = ", ierr
       call MPI_Finalize(ierr)
    endif
    yrms = sqrt(gsum/nx/ny)

    ierr = 0     ! return with success flag
    return
  end subroutine PRMS
  !---------------------------------------------------------------

end module UserData
!-----------------------------------------------------------------


!-----------------------------------------------------------------
! Main driver program
!-----------------------------------------------------------------
program driver

  ! inclusions
  use UserData
  implicit none
  include "mpif.h"

  ! Declarations
  ! general problem parameters
  real*8,    parameter :: pi = 3.1415926535897932d0
  real*8,    parameter :: T0 = 0.d0     ! initial time 
  real*8,    parameter :: Tf = 0.3d0    ! final time 
  integer,   parameter :: Nt = 20       ! total number of output times 
  real*8,    parameter :: rtol = 1.d-5  ! relative and absolute tolerances
  real*8,    parameter :: atol = 1.d-10
  integer*8, parameter :: nx_ = 60      ! spatial mesh size
  integer*8, parameter :: ny_ = 120
  real*8,    parameter :: kx_ = 0.5d0   ! heat conductivity coefficients
  real*8,    parameter :: ky_ = 0.75d0
  real*8,    parameter :: nlscoef = 1.d-7   ! nonlinear solver tolerance factor
  integer,   parameter :: PCGpretype = 1    ! enable preconditioner
  integer,   parameter :: PCGmaxl = 20      ! max num. PCG iterations
  real*8,    parameter :: PCGdelt = 0.d0    ! use default solver tolerance factor

  ! solution vector and other local variables
  real*8, allocatable :: y(:,:)
  real*8    :: rout(6), rpar, t, dTout, tout, urms
  integer*8 :: N, Ntot, i, j, iout(22), ipar
  integer   :: flag, ioutput
  logical   :: outproc
  character*100 :: outname

  ! initialize MPI
  call MPI_Init(flag)
  if (flag /= MPI_SUCCESS) then
     write(0,*) "Error in MPI_Init = ", flag
     stop
  end if
  call MPI_Comm_rank(MPI_COMM_WORLD, myid, flag)
  if (flag /= MPI_SUCCESS) then
     write(0,*) "Error in MPI_Comm_rank = ", flag
     call MPI_Finalize(flag)
  end if

  ! Initialize UserData module
  call InitUserData()
  nx = nx_
  ny = ny_
  kx = kx_
  ky = ky_
  dx = 1.d0/(nx-1)   ! x mesh spacing 
  dy = 1.d0/(ny-1)   ! x mesh spacing 

  ! Set up parallel decomposition (computes local mesh sizes)
  call SetupDecomp(flag)
  if (flag /= MPI_SUCCESS) then
     write(0,*) "Error in SetupDecomp = ", flag
     call MPI_Finalize(flag)
  end if

  ! Initial problem output 
  outproc = (myid == 0)
  if (outproc) then
     write(6,*) "  "
     write(6,*) "2D Heat PDE test problem:";
     write(6,'(A,i4)') "   nprocs = " , nprocs
     write(6,'(A,i4)') "   nx = ", nx
     write(6,'(A,i4)') "   ny = ", ny
     write(6,'(A,f5.2)') "   kx = ", kx
     write(6,'(A,f5.2)') "   ky = ", ky
     write(6,'(A,es9.2)') "   rtol = ", rtol
     write(6,'(A,es9.2)') "   atol = ", atol
     write(6,'(A,i4)') "   nxl (proc 0) = ", nxl
     write(6,'(A,i4)') "   nyl (proc 0) = ", nyl
     write(6,*) "  "
  endif

  ! Initialize data structures 
  N = nxl*nyl
  Ntot = nx*ny
  call FNVInitP(comm, 4, N, Ntot, flag)
  if (flag /= MPI_SUCCESS) then
     write(0,*) "Error in FNVInitP = ", flag
     call MPI_Finalize(flag)
  end if
  allocate(y(nxl,nyl))         ! Create parallel vector for solution 
  y = 0.d0                     ! Set initial conditions 

  ! Create the solver memory to use DIRK integrator, scalar tolerances
  call FARKMalloc(T0, y, 0, 1, rtol, atol, iout, rout, ipar, rpar, flag)
  if (flag /= MPI_SUCCESS) then
     write(0,*) "Error in FARKMalloc = ", flag
     call MPI_Finalize(flag)
  end if

  ! fill in the heat source array
  do j=1,nyl
     do i=1,nxl
        h(i,j) = sin(pi*(is+i-2)*dx) * sin(2.d0*pi*(js+j-2)*dy)
     enddo
  enddo

  ! set integrator options
  call FARKSetRin("NLCONV_COEF", nlscoef, flag)
  if (flag < 0) then
     write(0,*) "Error in FARKSetRin = ", flag
     call MPI_Finalize(flag)
  end if
  call FARKSetIin("PREDICT_METHOD", 1_8, flag)
  if (flag < 0) then
     write(0,*) "Error in FARKSetIin = ", flag
     call MPI_Finalize(flag)
  end if


  ! Linear solver specification 
  call FARKPCG(PCGpretype, PCGmaxl, PCGdelt, flag)  ! Specify the PCG solver 
  if (flag < 0) then
     write(0,*) "Error in FARKPCG = ", flag
     call MPI_Finalize(flag)
  end if
  call FARKSpilsSetPrec(1, flag)                     ! Enable preconditioning
  if (flag < 0) then
     write(0,*) "Error in FARKSpilsSetPrec = ", flag
     call MPI_Finalize(flag)
  end if

  ! specify that the problem is linearly implicit, but that Jacobian does not depend on time
  call FARKSetIin("LINEAR", 0_8, flag)
  if (flag < 0) then
     write(0,*) "Error in FARKSetIin = ", flag
     call MPI_Finalize(flag)
  end if

  ! Each processor outputs subdomain information
  write(outname,'(16Hheat2d_subdomain,f4.3,4H.txt)') myid/1000.0
  open(100, file=outname)
  write(100,'(6(i9,1x))') nx, ny, is, ie, js, je
  close(100)

  ! Open output streams for results, access data array 
  write(outname,'(6Hheat2d,f4.3,4H.txt)') myid/1000.0
  open(101, file=outname)

  ! Output initial condition to disk
  do j=1,nyl
     do i=1,nxl
        write(101,'(es25.16)',advance='no') y(i,j)
     enddo
  enddo
  write(101,*) "  "

  ! Main time-stepping loop: calls ARKode to perform the integration, then
  ! prints results.  Stops when the final time has been reached
  t = T0
  dTout = (Tf-T0)/Nt
  tout = T0+dTout
  call PRMS(y, urms, flag)
  if (outproc) then
    write(6,*) "        t      ||u||_rms"
    write(6,*) "   ----------------------"
    write(6,'(2(2x,f10.6))') t, urms
  endif
  do ioutput=1,Nt

     call FARKode(tout, t, y, 1, flag)         ! call integrator 
     if (flag < 0) then
        write(0,*) "Error in FARKode = ", flag
        exit
     end if
     
     call PRMS(y, urms, flag)
     if (outproc) &
          write(6,'(2(2x,f10.6))') t, urms     ! print solution stats 
     if (flag >= 0) then                       ! successful solve: update output time
        tout = min(tout + dTout, Tf)
     else                                      ! unsuccessful solve: break 
        if (outproc) &
             write(0,*) "Solver failure, stopping integration"
        exit
     endif

     ! output results to disk 
     do j=1,nyl
        do i=1,nxl
           write(101,'(es25.16)',advance='no') y(i,j)
        enddo
     enddo
     write(101,*) "  "
     
  enddo
  if (outproc) then
     write(6,*) "   ----------------------"
  endif
  close(101)

  ! Print some final statistics 
  if (outproc) then
     write(6,*) "  "
     write(6,*) "Final Solver Statistics:"
     write(6,'(2(A,i6),A)') "   Internal solver steps = ", iout(3), &
          " (attempted = ", iout(6), ")"
     write(6,'(A,i6,A,i6)') "   Total RHS evals:  Fe = ", iout(7), ",  Fi = ", iout(8)
     write(6,'(A,i6)') "   Total linear solver setups = ", iout(9)
     write(6,'(A,i6)') "   Total linear iterations = ", iout(21)
     write(6,'(A,i6)') "   Total number of Jacobian-vector products = ", iout(18)
     write(6,'(A,i6)') "   Total number of Preconditioner setups = ", iout(19)
     write(6,'(A,i6)') "   Total number of Preconditioner solves = ", iout(20)
     write(6,'(A,i6)') "   Total number of linear solver convergence failures = ", &
          iout(22)
     write(6,'(A,i6)') "   Total number of Newton iterations = ", iout(11)
     write(6,'(A,i6)') "   Total number of nonlinear solver convergence failures = ", &
          iout(12)
     write(6,'(A,i6)') "   Total number of error test failures = ", iout(10)
 endif

 ! Clean up and return with successful completion 
 deallocate(y)              ! free vectors
 deallocate(h)
 deallocate(d)
 call FreeUserData(flag)    ! free user data 
 call FARKFree()            ! free integrator memory 
 call MPI_Barrier(comm, flag)
 call MPI_Finalize(flag)    ! Finalize MPI

end program driver
!-----------------------------------------------------------------



!-----------------------------------------------------------------
! Functions called by the solver
!-----------------------------------------------------------------


subroutine farkifun(t, y, ydot, ipar, rpar, ierr)
!-----------------------------------------------------------------
! f routine to compute the ODE RHS function f(t,y). 
!-----------------------------------------------------------------
  ! declarations
  use UserData
  implicit none
  include "mpif.h"
  real*8,  intent(in)  :: t, rpar
  integer, intent(in)  :: ipar
  integer, intent(out) :: ierr
  real*8,  intent(in)  :: y(nxl,nyl)
  real*8,  intent(out) :: ydot(nxl,nyl)
  real*8    :: c1, c2, c3
  integer*8 :: i, j
  
  ! internals

  ! Initialize ydot to zero 
  ydot = 0.d0

  ! Exchange boundary data with neighbors
  call Exchange(y, ierr)
  if (ierr /= MPI_SUCCESS) then
     write(0,*) "Error in Exchange = " , ierr
     return
  end if

  ! iterate over subdomain interior, computing approximation to RHS
  c1 = kx/dx/dx
  c2 = ky/dy/dy
  c3 = -2.d0*(c1 + c2)
  do j=2,nyl-1
     do i=2,nxl-1
        ydot(i,j) = c1*(y(i-1,j)+y(i+1,j)) + c2*(y(i,j-1)+y(i,j+1)) + c3*y(i,j)
     enddo
  enddo
  
  ! iterate over subdomain boundaries (if not at overall domain boundary)
  if (.not. HaveBdry(1,1)) then    ! West face
     i=1
     do j=2,nyl-1
        ydot(i,j) = c1*(Wrecv(j)+y(i+1,j)) + c2*(y(i,j-1)+y(i,j+1)) + c3*y(i,j)
     enddo
  endif
  if (.not. HaveBdry(1,2)) then    ! East face
     i=nxl
     do j=2,nyl-1
        ydot(i,j) = c1*(y(i-1,j)+Erecv(j)) + c2*(y(i,j-1)+y(i,j+1)) + c3*y(i,j)
     enddo
  endif
  if (.not. HaveBdry(2,1)) then    ! South face
     j=1
     do i=2,nxl-1
        ydot(i,j) = c1*(y(i-1,j)+y(i+1,j)) + c2*(Srecv(i)+y(i,j+1)) + c3*y(i,j)
     enddo
  endif
  if (.not. HaveBdry(2,2)) then    ! West face
     j=nyl
     do i=2,nxl-1
        ydot(i,j) = c1*(y(i-1,j)+y(i+1,j)) + c2*(y(i,j-1)+Nrecv(i)) + c3*y(i,j)
     enddo
  endif
  if (.not. HaveBdry(1,1) .and. .not. HaveBdry(2,1)) then  ! South-West corner
     i=1
     j=1
     ydot(i,j) = c1*(Wrecv(j)+y(i+1,j)) + c2*(Srecv(i)+y(i,j+1)) + c3*y(i,j)
  endif
  if (.not. HaveBdry(1,1) .and. .not. HaveBdry(2,2)) then  ! North-West corner
     i=1
     j=nyl
     ydot(i,j) = c1*(Wrecv(j)+y(i+1,j)) + c2*(y(i,j-1)+Nrecv(i)) + c3*y(i,j)
  endif
  if (.not. HaveBdry(1,2) .and. .not. HaveBdry(2,1)) then  ! South-East corner
     i=nxl 
     j=1
     ydot(i,j) = c1*(y(i-1,j)+Erecv(j)) + c2*(Srecv(i)+y(i,j+1)) + c3*y(i,j)
  endif
  if (.not. HaveBdry(1,2) .and. .not. HaveBdry(2,2)) then  ! North-East corner
     i=nxl
     j=nyl
     ydot(i,j) = c1*(y(i-1,j)+Erecv(j)) + c2*(y(i,j-1)+Nrecv(i)) + c3*y(i,j)
  endif

  ydot = ydot + h         ! add in heat source

  ierr = 0                ! Return with success 
  return
end subroutine farkifun
!-----------------------------------------------------------------


subroutine farkefun(t, y, ydot, ipar, rpar, ierr)
!-----------------------------------------------------------------
! Explicit portion of f routine (empty)
!-----------------------------------------------------------------
  ! declarations
  use UserData
  implicit none
  real*8,  intent(in)  :: t, rpar
  integer, intent(in)  :: ipar
  integer, intent(out) :: ierr
  real*8,  intent(in)  :: y(nxl,nyl)
  real*8,  intent(out) :: ydot(nxl,nyl)
  
  ! internals

  ! Initialize ydot to zero and return with success
  ydot = 0.d0
  ierr = 0
  return
end subroutine farkefun
!-----------------------------------------------------------------


subroutine farkpset(t, y, fy, jok, jcur, gamma, hcur, ipar, &
                    rpar, v1, v2, v3, ierr)
!-----------------------------------------------------------------
! Preconditioner setup routine (fills inverse of Jacobian diagonal)
!-----------------------------------------------------------------
  ! declarations
  use UserData
  implicit none
  real*8,  intent(in)  :: t, gamma, hcur, rpar
  integer, intent(in)  :: ipar, jok
  integer, intent(out) :: jcur, ierr
  real*8,  intent(in)  :: y(nxl,nyl), fy(nxl,nyl)
  real*8               :: v1(nxl,nyl), v2(nxl,nyl), v3(nxl,nyl)
  real*8 :: c

  ! internals
  c = 1.d0 + gamma*2.d0*(kx/dx/dx + ky/dy/dy)

  ! set all entries of d to the inverse of the diagonal values in interior
  ! (since boundary RHS is 0, set boundary diagonals to the same)
  d = 1.d0/c

  jcur = 1     ! update jcur flag
  ierr = 0     ! Return with success 
  return
end subroutine farkpset


subroutine farkpsol(t, y, fy, r, z, gamma, delta, lr, &
                    ipar, rpar, vt, ierr)
!-----------------------------------------------------------------
! Preconditioner solve routine
!-----------------------------------------------------------------
  ! declarations
  use UserData
  implicit none
  real*8,  intent(in)  :: t, gamma, delta, rpar
  integer, intent(in)  :: lr, ipar
  integer, intent(out) :: ierr
  real*8,  intent(in)  :: y(nxl,nyl), fy(nxl,nyl), r(nxl,nyl)
  real*8,  intent(out) :: z(nxl,nyl)
  real*8               :: vt(nxl,nyl)

  ! internals
  z = r*d      ! perform Jacobi iteration (whole array operation)
  ierr = 0     ! Return with success 
  return
end subroutine farkpsol
!-----------------------------------------------------------------

