! ----------------------------------------------------------------
! Programmer(s): Daniel R. Reynolds @ SMU
!                Radu Serban and Alan C. Hindmarsh @ LLNL
!                modified by Daniel M. Margolis @ SMU
! ----------------------------------------------------------------
! SUNDIALS Copyright Start
! Copyright (c) 2002-2023, Lawrence Livermore National Security
! and Southern Methodist University.
! All rights reserved.
!
! See the top-level LICENSE and NOTICE files for details.
!
! SPDX-License-Identifier: BSD-3-Clause
! SUNDIALS Copyright End
! ----------------------------------------------------------------
! Example problem for FIDA: 2D heat equation, parallel, GMRES,
! IDABBDPRE.
!
! This example solves a discretized 2D heat equation problem.
! This version uses the Krylov solver SUNSPGMR and BBD
! preconditioning.
!
! The DAE system solved is a spatial discretization of the PDE
!          du/dt = d^2u/dx^2 + d^2u/dy^2
! on the unit square. The boundary condition is u = 0 on all edges.
! Initial conditions are given by u = 16 x (1 - x) y (1 - y). The
! PDE is treated with central differences on a uniform MX x MY
! grid. The values of u at the interior points satisfy ODEs, and
! equations u = 0 at the boundaries are appended, to form a DAE
! system of size N = MX * MY. Here MX = MY = 10.
!
! The system is actually implemented on submeshes, processor by
! processor, with an MXSUB by MYSUB mesh on each of NPEX * NPEY
! processors.
!
! The system is solved with FIDA using the Krylov linear solver
! SUNSPGMR in conjunction with the preconditioner module IDABBDPRE.
! The preconditioner uses a tridiagonal approximation
! (half-bandwidths = 1). The constraints u >= 0 are posed for all
! components. Local error testing on the boundary values is
! suppressed. Output is taken at t = 0, .01, .02, .04, ..., 10.24.
! ----------------------------------------------------------------

module Heat2DKryBBD_mod
  ! --------------------------------------------------------------
  ! 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.
  ! --------------------------------------------------------------

  !======= Inclusions ===========
  use, intrinsic :: iso_c_binding

  use fsundials_nvector_mod

  !======= Declarations =========
  implicit none

  ! With MPI-3 use mpi_f08 is preferred
  include "mpif.h"

  save

  ! SUNDIALS simulation context
  type(c_ptr) :: sunctx

  ! Grid and MPI domain decomposition information
  integer :: nx            ! global number of x grid points
  integer :: ny            ! global number of y grid points
  integer :: is            ! global x indices of this subdomain
  integer :: ie
  integer :: js            ! global y indices of this subdomain
  integer :: je
  integer :: nxl           ! local number of x grid points
  integer :: nyl           ! local number of y grid points
  integer(c_long) :: N, Ntot
  real(c_double)  :: dx    ! x-directional mesh spacing
  real(c_double)  :: dy    ! y-directional mesh spacing
  integer, target :: comm  ! communicator object
  integer :: myid          ! MPI process ID
  integer :: nprocs        ! total number of MPI processes
  logical :: HaveNbor(2,2) ! flags denoting neighbor on boundary
  real(c_double), dimension(:), allocatable :: Erecv  ! receive buffers for neighbor exchange
  real(c_double), dimension(:), allocatable :: Wrecv
  real(c_double), dimension(:), allocatable :: Nrecv
  real(c_double), dimension(:), allocatable :: Srecv
  real(c_double), dimension(:), allocatable :: Esend  ! send buffers for neighbor exchange
  real(c_double), dimension(:), allocatable :: Wsend
  real(c_double), dimension(:), allocatable :: Nsend
  real(c_double), dimension(:), allocatable :: Ssend

  ! Problem parameters
  integer(c_long)  :: mudq, mldq, mu, ml
  integer(c_int)   :: maxl
  real(c_double) :: kx   ! x-directional diffusion coefficient
  real(c_double) :: ky   ! y-directional diffusion coefficient

  ! Printed parameters
  real(c_double)  :: h(1)       ! size of last step
  integer(c_int)  :: k(1)       ! number of last order
  integer(c_long) :: nst(1)     ! number of time steps
  integer(c_long) :: nre(1)     ! number of residual evals
  integer(c_long) :: nreLS(1)   ! number of residual linear solver evals
  integer(c_long) :: nge(1)     ! number of implicit RHS evals
  integer(c_long) :: netf(1)    ! number of error test fails
  integer(c_long) :: nncf(1)    ! number of nonlinear convergence fails
  integer(c_long) :: nlcf(1)    ! number of linear convergence fails
  integer(c_long) :: nni(1)     ! number of nonlinear iters
  integer(c_long) :: nli(1)     ! number of linear iters
  integer(c_long) :: npre(1)    ! number of preconditioner setups
  integer(c_long) :: npsol(1)   ! number of preconditioner solves

contains

  ! --------------------------------------------------------------
  ! Initialize memory allocated within Userdata (set to defaults)
  ! --------------------------------------------------------------
  subroutine InitHeat2DData()
    implicit none
    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
    comm = MPI_COMM_WORLD
    myid = 0
    nprocs = 0
    HaveNbor = .false.
    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 InitHeat2DData
  ! --------------------------------------------------------------


  ! --------------------------------------------------------------
  ! Set up parallel decomposition
  ! --------------------------------------------------------------
  subroutine SetupDecomp(ierr)
    ! declarations
    implicit none

    integer, intent(out) :: ierr
    integer :: dims(2), periods(2), coords(2)

    ! internals

    ! 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
    HaveNbor(1,1) = (is /= 1)
    HaveNbor(1,2) = (ie /= nx)
    HaveNbor(2,1) = (js /= 1)
    HaveNbor(2,2) = (je /= ny)
    if (HaveNbor(1,1)) then
       allocate(Wrecv(nyl))
       allocate(Wsend(nyl))
    end if
    if (HaveNbor(1,2)) then
       allocate(Erecv(nyl))
       allocate(Esend(nyl))
    end if
    if (HaveNbor(2,1)) then
       allocate(Srecv(nxl))
       allocate(Ssend(nxl))
    end if
    if (HaveNbor(2,2)) then
       allocate(Nrecv(nxl))
       allocate(Nsend(nxl))
    end if

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

  ! --------------------------------------------------------------
  ! Free memory allocated within Userdata
  ! --------------------------------------------------------------
  subroutine FreeHeat2DData(ierr)
    implicit none
    integer, intent(out) :: ierr
    if (allocated(Wrecv))  deallocate(Wrecv)
    if (allocated(Wsend))  deallocate(Wsend)
    if (allocated(Erecv))  deallocate(Erecv)
    if (allocated(Esend))  deallocate(Esend)
    if (allocated(Srecv))  deallocate(Srecv)
    if (allocated(Ssend))  deallocate(Ssend)
    if (allocated(Nrecv))  deallocate(Nrecv)
    if (allocated(Nsend))  deallocate(Nsend)
    ierr = 0     ! return with success flag
    return
  end subroutine FreeHeat2DData
  ! --------------------------------------------------------------

  subroutine InitProfile(sunvec_y, sunvec_ydot, sunvec_id, &
       sunvec_res, sunvec_c, ierr)
    use fnvector_parallel_mod
    implicit none
    type(N_Vector), pointer, intent(inout) :: sunvec_y
    type(N_Vector), pointer, intent(inout) :: sunvec_ydot
    type(N_Vector), pointer, intent(inout) :: sunvec_id
    type(N_Vector), pointer, intent(inout) :: sunvec_res
    type(N_Vector), pointer, intent(inout) :: sunvec_c
    integer(c_int), intent(in) :: ierr
    real(c_double), pointer, dimension(nxl,nyl) :: y(:,:), ydot(:,:), id(:,:), res(:,:), cstr(:,:)
    real(c_double) :: xreal, yreal
    integer(c_int) :: retval
    type(c_ptr), value :: user_data
    integer :: i, j
    ! Create solution vector, point at its data, and set initial condition
    N = nxl*nyl
    Ntot = nx*ny
    sunvec_y    => FN_VNew_Parallel(comm, N, Ntot, sunctx)
    sunvec_ydot => FN_VNew_Parallel(comm, N, Ntot, sunctx)
    sunvec_id   => FN_VNew_Parallel(comm, N, Ntot, sunctx)
    sunvec_res  => FN_VNew_Parallel(comm, N, Ntot, sunctx)
    sunvec_c    => FN_VNew_Parallel(comm, N, Ntot, sunctx)
    y(1:nxl,1:nyl)    => FN_VGetArrayPointer(sunvec_y)
    ydot(1:nxl,1:nyl) => FN_VGetArrayPointer(sunvec_ydot)
    id(1:nxl,1:nyl)   => FN_VGetArrayPointer(sunvec_id)
    res(1:nxl,1:nyl)  => FN_VGetArrayPointer(sunvec_res)
    cstr(1:nxl,1:nyl) => FN_VGetArrayPointer(sunvec_c)
    id = 1.d0
    do i = 1,nxl
       xreal = dx*dble(is+i-2)
       do j = 1,nyl
          yreal = dy*dble(js+j-2)
          if (.not. HaveNbor(1,1) .and. i == 1) then
             id(i,j) = 0.d0
          end if
          if (.not. HaveNbor(1,2) .and. i == nxl) then
             id(i,j) = 0.d0
          end if
          if (.not. HaveNbor(2,1) .and. j == 1) then
             id(i,j) = 0.d0
          end if
          if (.not. HaveNbor(2,2) .and. j == nyl) then
             id(i,j) = 0.d0
          end if
          y(i,j) = 16.d0*xreal*(1.d0-xreal)*yreal*(1.d0-yreal)
       end do
    end do
    ydot = 0.d0
    cstr = 1.d0
    retval = resfn(0.d0, sunvec_y, sunvec_ydot, sunvec_res, user_data)
    if (retval /= 0) then
       print *, "Error: resfn in InitProfile returned ",retval
       call MPI_Abort(comm, 1, ierr)
    end if
    ydot = -1.d0*res

    return
  end subroutine InitProfile
  ! --------------------------------------------------------------

  subroutine getStats(ida_mem, retval, ierr)

    !======= Inclusions ===========
    use, intrinsic :: iso_c_binding
    use fida_mod
    use fsundials_futils_mod

    !======= Declarations =========
    implicit none

    ! calling variables
    type(c_ptr),    intent(in)  :: ida_mem
    integer(c_int), intent(in)  :: ierr
    integer(c_int), intent(out) :: retval

    retval = FIDAGetLastOrder(ida_mem, k)
    if (retval /= 0) then
       print *, "Error: FIDAGetLastOrder returned ",retval
       call MPI_Abort(comm, 1, ierr)
    end if

    retval = FIDAGetNumSteps(ida_mem, nst)
    if (retval /= 0) then
       print *, "Error: FIDAGetNumSteps returned ",retval
       call MPI_Abort(comm, 1, ierr)
    end if

    retval = FIDAGetNumNonlinSolvIters(ida_mem, nni)
    if (retval /= 0) then
       print *, "Error: FIDAGetNumNonlinSolvIters returned ",retval
       call MPI_Abort(comm, 1, ierr)
    end if

    retval = FIDAGetNumLinIters(ida_mem, nli)
    if (retval /= 0) then
       print *, "Error: FIDAGetNumLinIters returned ",retval
       call MPI_Abort(comm, 1, ierr)
    end if

    retval = FIDAGetNumResEvals(ida_mem, nre)
    if (retval /= 0) then
       print *, "Error: FIDAGetNumResEvals returned ",retval
       call MPI_Abort(comm, 1, ierr)
    end if

    retval = FIDAGetNumLinResEvals(ida_mem, nreLS)
    if (retval /= 0) then
       print *, "Error: FIDAGetNumLinResEvals returned ",retval
       call MPI_Abort(comm, 1, ierr)
    end if

    retval = FIDABBDPrecGetNumGfnEvals(ida_mem, nge)
    if (retval /= 0) then
       print *, "Error: FIDABBDPrecGetNumGfnEvals returned ",retval
       call MPI_Abort(comm, 1, ierr)
    end if

    retval = FIDAGetLastStep(ida_mem, h)
    if (retval /= 0) then
       print *, "Error: FIDAGetLastStep returned ",retval
       call MPI_Abort(comm, 1, ierr)
    end if

    retval = FIDAGetNumPrecEvals(ida_mem, npre)
    if (retval /= 0) then
       print *, "Error: FIDAGetNumPrecEvals returned ",retval
       call MPI_Abort(comm, 1, ierr)
    end if

    retval = FIDAGetNumPrecSolves(ida_mem, npsol)
    if (retval /= 0) then
       print *, "Error: FIDAGetNumPrecSolves returned ",retval
       call MPI_Abort(comm, 1, ierr)
    end if

  end subroutine getStats
  ! --------------------------------------------------------------

  ! ----------------------------------------------------------------
  ! DAE residual function r(t,y).
  ! ----------------------------------------------------------------
  integer(c_int) function resfn(t, sunvec_y, sunvec_ydot, sunvec_res, &
       user_data) result(retval) bind(C)

    !======= Inclusions ===========
    use, intrinsic :: iso_c_binding
    use fsundials_nvector_mod

    !======= Declarations =========
    implicit none

    ! calling variables
    real(c_double), value :: t            ! current time
    type(N_Vector)        :: sunvec_y     ! solution N_Vector
    type(N_Vector)        :: sunvec_ydot  ! rhs N_Vector
    type(N_Vector)        :: sunvec_res   ! residual N_Vector
    type(c_ptr),    value :: user_data    ! user-defined data

    !======= Internals ============

    ! Exchange boundary data with neighbors
    retval = Exchange(N, t, sunvec_y, sunvec_ydot, sunvec_res, user_data)
    if (retval /= MPI_SUCCESS) then
       write(0,*) "Error in Exchange = " , retval
       return
    end if

    retval = LocalFn(N, t, sunvec_y, sunvec_ydot, sunvec_res, user_data)
    if (retval /= MPI_SUCCESS) then
       write(0,*) "Error in LocalFn = " , retval
       return
    end if

    retval = 0              ! Return with success
    return
  end function resfn
  ! ----------------------------------------------------------------


  ! ----------------------------------------------------------------
  ! Perform neighbor exchange (Communication function)
  ! ----------------------------------------------------------------
  integer(c_int) function Exchange(Nloc, t, sunvec_y, sunvec_ydot, &
       sunvec_g, user_data) result(ierr) bind(C)

    !======= Inclusions ===========
    use, intrinsic :: iso_c_binding
    use fsundials_nvector_mod

    !======= Declarations =========
    implicit none

    ! calling variables
    integer(c_long), value :: Nloc
    real(c_double),  value :: t            ! current time
    type(N_Vector)         :: sunvec_y     ! solution N_Vector
    type(N_Vector)         :: sunvec_ydot  ! rhs N_Vector
    type(N_Vector)         :: sunvec_g     ! evaluated N_Vector
    type(c_ptr),     value :: user_data    ! user-defined data

    real(c_double), pointer, dimension(nxl,nyl) :: y(:,:)
    integer :: reqSW, reqSE, reqSS, reqSN, reqRW, reqRE, reqRS, reqRN;
    integer :: stat(MPI_STATUS_SIZE)
    integer :: i, ipW, ipE, ipS, ipN
    integer :: coords(2), dims(2), periods(2), nbcoords(2)

    ! internals
    y(1:nxl,1:nyl) => FN_VGetArrayPointer(sunvec_y)

    ! 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
    end if
    if (HaveNbor(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
       end if
    end if
    if (HaveNbor(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
       end if
    end if
    if (HaveNbor(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
       end if
    end if
    if (HaveNbor(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
       end if
    end if

    ! open Irecv buffers
    if (HaveNbor(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
       end if
    end if
    if (HaveNbor(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
       end if
    end if
    if (HaveNbor(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
       end if
    end if
    if (HaveNbor(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
       end if
    end if

    ! send data
    if (HaveNbor(1,1)) then
       do i=1,nyl
          Wsend(i) = y(1,i)
       end do
       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
       end if
    end if
    if (HaveNbor(1,2)) then
       do i=1,nyl
          Esend(i) = y(nxl,i)
       end do
       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
       end if
    end if
    if (HaveNbor(2,1)) then
       do i=1,nxl
          Ssend(i) = y(i,1)
       end do
       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
       end if
    end if
    if (HaveNbor(2,2)) then
       do i=1,nxl
          Nsend(i) = y(i,nyl)
       end do
       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
       end if
    end if

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

    ierr = MPI_SUCCESS    ! return with success flag
    return
  end function Exchange
  ! ----------------------------------------------------------------


  ! ----------------------------------------------------------------
  ! Processor-local portion of the DAE residual function.
  ! ----------------------------------------------------------------
  integer(c_int) function LocalFn(Nloc, t, sunvec_y, sunvec_ydot, sunvec_g, &
       user_data) result(ierr) bind(C)

    !======= Inclusions ===========
    use, intrinsic :: iso_c_binding
    use fsundials_nvector_mod

    !======= Declarations =========
    implicit none

    ! calling variables
    integer(c_long), value :: Nloc
    real(c_double),  value :: t            ! current time
    type(N_Vector)         :: sunvec_y     ! solution N_Vector
    type(N_Vector)         :: sunvec_ydot  ! rhs N_Vector
    type(N_Vector)         :: sunvec_g     ! evaluated N_Vector
    type(c_ptr),     value :: user_data    ! user-defined data

    ! pointers to data in SUNDIALS vectors
    real(c_double), pointer, dimension(nxl,nyl) :: y(:,:)
    real(c_double), pointer, dimension(nxl,nyl) :: ydot(:,:)
    real(c_double), pointer, dimension(nxl,nyl) :: res(:,:)

    ! local data
    real(c_double) :: c1, c2, c3
    integer :: i, j

    !======= Internals ============

    ! Get data arrays from SUNDIALS vectors
    y(1:nxl,1:nyl)    => FN_VGetArrayPointer(sunvec_y)
    ydot(1:nxl,1:nyl) => FN_VGetArrayPointer(sunvec_ydot)
    res(1:nxl,1:nyl)  => FN_VGetArrayPointer(sunvec_g)

    ! set constants
    c1 = kx/dx/dx
    c2 = ky/dy/dy
    c3 = -2.d0*(c1 + c2)

    ! Copy y NVector into res NVector
    res = y

    ! iterate over subdomain boundaries (if not at overall domain boundary)
    do i = 1,nxl
       do j = 1,nyl
          if (i == 1 .and. j == 1) then
             if (HaveNbor(1,1) .and. HaveNbor(2,1)) then  ! South-West corner
                res(i,j) = c1*(Wrecv(j)+y(i+1,j)) + c2*(Srecv(i)+y(i,j+1)) + c3*y(i,j)
             end if
          else if (i == 1 .and. j == nyl) then
             if (HaveNbor(1,1) .and. HaveNbor(2,2)) then  ! North-West corner
                res(i,j) = c1*(Wrecv(j)+y(i+1,j)) + c2*(y(i,j-1)+Nrecv(i)) + c3*y(i,j)
             end if
          else if (i == nxl .and. j == 1) then
             if (HaveNbor(1,2) .and. HaveNbor(2,1)) then  ! South-East corner
                res(i,j) = c1*(y(i-1,j)+Erecv(j)) + c2*(Srecv(i)+y(i,j+1)) + c3*y(i,j)
             end if
          else if (i == nxl .and. j == nyl) then
             if (HaveNbor(1,2) .and. HaveNbor(2,2)) then  ! North-East corner
                res(i,j) = c1*(y(i-1,j)+Erecv(j)) + c2*(y(i,j-1)+Nrecv(i)) + c3*y(i,j)
             end if
          else if (i == 1) then
             if (HaveNbor(1,1)) then                      ! West face
                res(i,j) = c1*(Wrecv(j)+y(i+1,j)) + c2*(y(i,j-1)+y(i,j+1)) + c3*y(i,j)
             end if
          else if (i == nxl) then
             if (HaveNbor(1,2)) then                      ! East face
                res(i,j) = c1*(y(i-1,j)+Erecv(j)) + c2*(y(i,j-1)+y(i,j+1)) + c3*y(i,j)
             end if
          else if (j == 1) then
             if (HaveNbor(2,1)) then                      ! South face
                res(i,j) = c1*(y(i-1,j)+y(i+1,j)) + c2*(Srecv(i)+y(i,j+1)) + c3*y(i,j)
             end if
          else if (j == nyl) then
             if (HaveNbor(2,2)) then                      ! North face
                res(i,j) = c1*(y(i-1,j)+y(i+1,j)) + c2*(y(i,j-1)+Nrecv(i)) + c3*y(i,j)
             end if
          else
             res(i,j) = c1*(y(i-1,j)+y(i+1,j)) + c2*(y(i,j-1)+y(i,j+1)) + c3*y(i,j)
          end if
          res(i,j) = ydot(i,j) - res(i,j)
       end do
    end do

    ierr = 0     ! Return with success
    return
  end function LocalFn
  ! ----------------------------------------------------------------

end module Heat2DKryBBD_mod
! ------------------------------------------------------------------


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

  ! inclusions
  use, intrinsic :: iso_c_binding
  use fsundials_futils_mod       ! Fortran utilities
  use fida_mod                   ! Access IDA
  use fsundials_types_mod        ! sundials defined types
  use fsundials_nvector_mod      ! Access generic N_Vector
  use fnvector_parallel_mod      ! Access parallel N_Vector
  use fsundials_context_mod      ! Access sundials context
  use fsundials_matrix_mod       ! Access generic SUNMatrix
  use fsundials_linearsolver_mod ! Access generic SUNLinearSolver
  use fsunlinsol_spgmr_mod       ! Access GMRES SUNLinearSolver

  use Heat2DKryBBD_mod

  !======= Declarations =========
  implicit none

  ! Declarations
  ! general problem parameters
  integer, parameter :: Nt = 11           ! total number of output times
  integer, parameter :: nx_ = 100         ! spatial mesh size
  integer, parameter :: ny_ = 100
  real(c_double), parameter :: T0 = 0.d0        ! initial time
  real(c_double), parameter :: T1 = 0.01d0      ! final time
  real(c_double), parameter :: rtol = 0.d0      ! relative and absolute tolerances
  real(c_double), parameter :: atol = 1.d-3
  real(c_double), parameter :: kx_ = 1.0d0      ! heat conductivity coefficients
  real(c_double), parameter :: ky_ = 1.0d0

  ! solution vector and other local variables
  type(N_Vector), pointer :: sunvec_y        ! solution N_Vector
  type(N_Vector), pointer :: sunvec_f        ! derivative N_Vector
  type(N_Vector), pointer :: sunvec_id       ! derivative N_Vector
  type(N_Vector), pointer :: sunvec_res      ! derivative N_Vector
  type(N_Vector), pointer :: sunvec_c        ! constraint N_Vector
  real(c_double), pointer, dimension(nxl,nyl) :: y(:,:)  ! vector data
  real(c_double), pointer, dimension(nxl,nyl) :: f(:,:)  ! vector data
  type(SUNLinearSolver), pointer :: sun_LS   ! linear solver
  type(SUNMatrix),       pointer :: sunmat_A ! sundials matrix
  type(c_ptr)     :: ida_mem                 ! IDA memory
  integer(c_int) :: retval
  integer :: ierr, case
  logical :: outproc
  real(c_double) :: t(1), dTout, tout, ymax
  integer :: i, j, ioutput
  character*100 :: outname

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

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

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

  ! Create SUNDIALS simulation context, now that comm has been configured
  retval = FSUNContext_Create(comm, sunctx)
  if (retval /= 0) then
     print *, "Error: FSUNContext_Create returned ",retval
     call MPI_Abort(comm, 1, ierr)
  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,*) "  "
  end if

  ! Create the IDA timestepper module
  ida_mem = FIDACreate(sunctx)
  if (.not. c_associated(ida_mem)) then
     print *, "Error: FIDACreate returned NULL"
     call MPI_Abort(comm, 1, ierr)
  end if

  call InitProfile(sunvec_y, sunvec_f, sunvec_id, sunvec_res, sunvec_c, ierr)

  retval = FIDAInit(ida_mem, c_funloc(resfn), t0, sunvec_y, sunvec_f)
  if (retval /= 0) then
     print *, "Error: FIDAInit returned ",retval
     call MPI_Abort(comm, 1, ierr)
  end if

  ! Create linear solver
  maxl = 0
  sun_LS => FSUNLinSol_SPGMR(sunvec_y, SUN_PREC_LEFT, maxl, sunctx)
  if (.not. associated(sun_LS)) then
     print *, "Error: FSUNLinSol_PCG returned NULL"
     call MPI_Abort(comm, 1, ierr)
  end if

  ! Attach linear solver
  sunmat_A => null()
  retval = FIDASetLinearSolver(ida_mem, sun_LS, sunmat_A)
  if (retval /= 0) then
     print *, "Error: FIDASetLinearSolver returned ",retval
     call MPI_Abort(comm, 1, ierr)
  end if

  retval = FSUNLinSol_SPGMRSetMaxRestarts(sun_LS, 5)
  if (retval /= 0) then
     print *, "Error: FSUNLinSol_SPGMRSetMaxRestarts returned",retval
     call MPI_Abort(comm, 1, ierr)
  end if

  ! Attach preconditioner
  mudq = nxl
  mldq = nxl
  mu = 1
  ml = 1
  retval = FIDABBDPrecInit(ida_mem, N, mudq, mldq, mu, ml, 0.d0, &
       c_funloc(LocalFn), c_funloc(Exchange))
  if (retval /= 0) then
     print *, "Error: FIDASetPreconditioner returned ",retval
     call MPI_Abort(comm, 1, ierr)
  end if

  ! Specify tolerances
  retval = FIDASStolerances(ida_mem, rtol, atol)
  if (retval /= 0) then
     print *, "Error: FIDASStolerances returned ",retval
     call MPI_Abort(comm, 1, ierr)
  end if

  retval = FIDASetSuppressAlg(ida_mem, SUNTRUE)
  if (retval /= 0) then
     print *, "Error: FIDASetSuppressAlg returned ",retval
     call MPI_Abort(comm, 1, ierr)
  end if

  retval = FIDASetId(ida_mem, sunvec_id)
  if (retval /= 0) then
     print *, "Error: FIDASetId returned ",retval
     call MPI_Abort(comm, 1, ierr)
  end if

  retval = FIDASetConstraints(ida_mem, sunvec_c)
  if (retval /= 0) then
     print *, "Error: FIDASetConstraints returned ",retval
     call MPI_Abort(comm, 1, ierr)
  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
  y(1:nxl,1:nyl) => FN_VGetArrayPointer(sunvec_y)
  do j=1,nyl
     do i=1,nxl
        write(101,'(es25.16)',advance='no') y(i,j)
     end do
  end do
  write(101,*) "  "

  do case = 1,2
     if (case == 2) then
        mudq = 5
        mldq = 5
        call InitProfile(sunvec_y, sunvec_f, sunvec_id, sunvec_res, sunvec_c, ierr)
        retval = FIDAReInit(ida_mem, t0, sunvec_y, sunvec_f)
        if (retval /= 0) then
           print *, "Error: FIDAReInit returned ",retval
           call MPI_Abort(comm, 1, ierr)
        end if
        retval = FIDABBDPrecReInit(ida_mem, mudq, mldq, 0.d0)
        if (retval /= 0) then
           print *, "Error: FIDABBDPrecReInit returned ",retval
           call MPI_Abort(comm, 1, ierr)
        end if
        if(outproc) then
           write(6,*) "  "
           write(6,*) "Case ", case
           write(6,*) "  Difference quotient half-bandwidths = ", mudq
           write(6,*) "  Retained matrix half-bandwidths = ", mu
           write(6,*) "  "
           write(6,*) "Output Summary"
           write(6,*) "  umax = max-norm of solution"
           write(6,*) "  nre = nre + nreLS (total number of RES evals.)"
        end if
     end if
     if (case == 1) then
        if(outproc) then
           write(6,*) "  "
           write(6,*) "Case ", case
           write(6,*) "  Difference quotient half-bandwidths = ", mudq
           write(6,*) "  Retained matrix half-bandwidths = ", mu
           write(6,*) "  "
           write(6,*) "Output Summary"
           write(6,*) "  umax = max-norm of solution"
           write(6,*) "  nre = nre + nreLS (total number of RES evals.)"
        end if
     end if
     ! Main time-stepping loop: calls IDA to perform the integration, then
     ! prints results.  Stops when the final time has been reached
     t = T0
     tout = T1
     if (outproc) then
        write(6,*) "   "
        write(6,*) "        t      ||u||_max       k  nst   nni   nli   nre    nge       h      npe   nps"
        write(6,*) "   ------------------------------------------------------------------------------------"
     end if
     do ioutput=1,Nt

        ! Integrate to output time
        retval = FIDASolve(ida_mem, tout, t, sunvec_y, sunvec_f, IDA_NORMAL)
        if (retval /= 0) then
           print *, "Error: FIDASolve returned ",retval
           call MPI_Abort(comm, 1, ierr)
        end if

        ! print solution stats and update internal time
        ymax = FN_VMaxNorm(sunvec_y)
        call getStats(ida_mem, retval, ierr)
        if (outproc)        write(6,'(2x,f10.6,2x,es13.5,3x,i1,3x,i2,3x,i3,3x,i3,3x,i3,a,i3,3x,i4,3x,es9.2,3x,i2,3x,i3)') &
             t, ymax, k, nst, nni, nli, nre, "+", nreLS, nge, h, npre, npsol
        tout = 2.0d0 * tout

        ! output results to disk
        do j=1,nyl
           do i=1,nxl
              write(101,'(es25.16)',advance='no') y(i,j)
           end do
        end do
        write(101,*) "  "

     end do
     if (outproc) then
        write(6,*) "   ------------------------------------------------------------------------------------"
     end if
     close(101)

     retval = FIDAGetNumErrTestFails(ida_mem, netf)
     if (retval /= 0) then
        print *, "Error: FIDAGetNumErrTestFails returned ",retval
        call MPI_Abort(comm, 1, ierr)
     end if

     retval = FIDAGetNumNonlinSolvConvFails(ida_mem, nncf)
     if (retval /= 0) then
        print *, "Error: FIDAInit returned ",retval
        call MPI_Abort(comm, 1, ierr)
     end if

     retval = FIDAGetNumLinConvFails(ida_mem, nlcf)
     if (retval /= 0) then
        print *, "Error: FIDAInit returned ",retval
        call MPI_Abort(comm, 1, ierr)
     end if

     ! Print some final statistics
     if (outproc) then
        write(6,*) "  "
        write(6,*) "Final Solver Statistics:"
        write(6,'(A,i6)') "   Total number of error test failures      = ", netf
        write(6,'(A,i6)') "   Total number of nonlinear conv. failures = ", nncf
        write(6,'(A,i6)') "   Total number of linear conv. failures    = ", nlcf
     end if
  end do

  ! Clean up and return with successful completion
  call FIDAFree(ida_mem)           ! free integrator memory
  call FN_VDestroy_Parallel(sunvec_y)          ! free vector memory
  call FN_VDestroy_Parallel(sunvec_f)
  call FN_VDestroy_Parallel(sunvec_id)
  call FN_VDestroy_Parallel(sunvec_res)
  call FN_VDestroy_Parallel(sunvec_c)
  retval = FSUNLinSolFree(sun_LS)     ! free linear solver
  call FreeHeat2DData(ierr)           ! free user data
  call MPI_Barrier(comm, ierr)
  call MPI_Finalize(ierr)             ! Finalize MPI

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