!>
!> @file test_jacobi.f90
!>
!> @brief
!>
!> @copyright
!> Copyright (©) 2021 EPFL (Ecole Polytechnique Fédérale de Lausanne)
!> SPC (Swiss Plasma Center)
!>
!> SPClibs is free software: you can redistribute it and/or modify it under
!> the terms of the GNU Lesser General Public License as published by the Free
!> Software Foundation, either version 3 of the License, or (at your option)
!> any later version.
!>
!> SPClibs is distributed in the hope that it will be useful, but WITHOUT ANY
!> WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS
!> FOR A PARTICULAR PURPOSE. See the GNU General Public License for more details.
!>
!> You should have received a copy of the GNU Lesser General Public License
!> along with this program. If not, see <https://www.gnu.org/licenses/>.
!>
!> @author
!> (in alphabetical order)
!> @author Trach-Minh Tran <trach-minh.tran@epfl.ch>
!>
!
!    Test 2D  parallel Jacobi using STENCIL_2D matrix-free structure.
!
MODULE mod
  USE iso_fortran_env, ONLY : rkind => real64
  IMPLICIT NONE
!
  LOGICAL, PARAMETER     :: nldebug=.FALSE.
  REAL(rkind), PARAMETER :: PI=4.d0*ATAN(1.0d0)
CONTAINS
END MODULE mod
!++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
PROGRAM main
  USE mpi
  USE pputils2, ONLY : dist1d, exchange, norm2_vec=>ppnorm2, timera, hostlist
  USE parmg, ONLY    : jacobi, get_resids
  USE stencil, ONLY  : stencil_2d, init, laplacian, putmat
  USE mod
  IMPLICIT NONE
!
  INTEGER, PARAMETER :: ndims=2
!
  INTEGER                   :: me, neighs(4), npes, ierr
  INTEGER, DIMENSION(ndims) :: dims=[0,0]
  INTEGER, DIMENSION(ndims) :: coords, comm1d
  LOGICAL, DIMENSION(ndims) :: periods=[.FALSE.,.FALSE.]
  LOGICAL                   :: reorder =.FALSE.
  INTEGER                   :: comm_cart
!
  INTEGER                   :: nx=4, ny=4    ! Number of intervals
  INTEGER, DIMENSION(ndims) :: e, s, lb, ub, npt_glob, npt_loc
!
  REAL(rkind), ALLOCATABLE  :: xgrid(:), ygrid(:)
  REAL(rkind)               :: dx, dy
  INTEGER, DIMENSION(5,2)   :: id     ! 5-point stencil 
  INTEGER                   :: npoints
  TYPE(stencil_2d)          :: mat
  INTEGER                   :: i
!
  REAL(rkind), ALLOCATABLE  :: f(:,:), v(:,:), u(:,:)
  REAL(rkind), ALLOCATABLE  :: resids(:,:), errs(:,:)
  REAL(rkind), ALLOCATABLE  :: resid_it(:), err_it(:)
  REAL(rkind)               :: omega=1.0d0, resid
  INTEGER                   :: it, it_skip, nits=100
!
  NAMELIST /in/ nx, ny, omega, nits
!================================================================================
!                        1.0  Prologue
!
!  2D process grid
  CALL mpi_init(ierr)
  CALL mpi_comm_rank(MPI_COMM_WORLD, me, ierr)
  CALL mpi_comm_size(MPI_COMM_WORLD, npes, ierr)
  CALL mpi_dims_create(npes, ndims, dims, ierr)
  CALL mpi_cart_create(MPI_COMM_WORLD, ndims, dims, periods, reorder, comm_cart,&
       &               ierr)
!
  CALL mpi_comm_rank(comm_cart, me, ierr)
  CALL mpi_cart_coords(comm_cart, me, ndims, coords, ierr)
  CALL mpi_cart_shift(comm_cart, 0, 1, neighs(1), neighs(2), ierr)
  CALL mpi_cart_shift(comm_cart, 1, 1, neighs(3), neighs(4), ierr)
!
  CALL mpi_cart_sub(comm_cart, [.TRUE.,.FALSE.], comm1d(1), ierr)
  CALL mpi_cart_sub(comm_cart, [.FALSE.,.TRUE.], comm1d(2), ierr)
!
  CALL hostlist(comm_cart)
  IF(me.EQ.0) THEN
     WRITE(*,'(a,i0,a,i0/)') "Process grid: ", dims(1), " X ", dims(2)
  END IF
!
!  Read problem inputs
  IF(me.EQ.0) THEN
     READ(*,in)
     WRITE(*,in)
  END IF
!
  CALL mpi_bcast(nx, 1, MPI_INTEGER, 0, comm_cart, ierr)
  CALL mpi_bcast(ny, 1, MPI_INTEGER, 0, comm_cart, ierr)
  CALL mpi_bcast(nits, 1, MPI_INTEGER, 0, comm_cart, ierr)
  CALL mpi_bcast(omega, 1, MPI_DOUBLE_PRECISION, 0, comm_cart, ierr)
!================================================================================
!                        2.0  2d Grid construction
!
!  Partition 2D grid
  CALL timera(0, 'Grid_construction')
  npt_glob(1) = nx+1
  npt_glob(2) = ny+1
  DO i=1,ndims
     CALL dist1d(comm1d(i), 0, npt_glob(i), s(i), npt_loc(i))
     e(i) = s(i) + npt_loc(i) - 1
  END DO
  WRITE(*,'(a,i3.3,a,2(3i4,", "))') 'PE', me, ' coords, s,e:', &
       &  (coords(i),s(i),e(i),i=1,ndims)
!
!   Global mesh
  dx = 1.0d0/REAL(nx)
  dy = 1.0d0/REAL(ny)
  ALLOCATE(xgrid(0:nx))
  ALLOCATE(ygrid(0:ny))
  xgrid = [ (i*dx, i=0,nx) ]
  ygrid = [ (i*dy, i=0,ny) ]
  CALL timera(1, 'Grid_construction')
!================================================================================
!                        3.0  FD Laplacian
!
  CALL timera(0, 'Laplacian')
  id=RESHAPE([ 0, -1, 0, 1, 0,  & 
               0,  0,-1, 0, 1], &
              [5,2])
  npoints = 5
  CALL init(s, e, id, .FALSE., mat, comm_cart)
!
  CALL laplacian(dx, dy, mat)
  CALL timera(1, 'Laplacian')
!================================================================================
!                        4.0 Test parallel Jacobi with \nabla u(x,y) = f(x,y)
!
!   Problem definition
!
  s = mat%s
  e = mat%e
  lb = s-1
  ub = e+1
  ALLOCATE(f(lb(1):ub(1),lb(2):ub(2)))       ! RHS
  ALLOCATE(v(lb(1):ub(1),lb(2):ub(2)))       ! Exact solutions
  ALLOCATE(u(lb(1):ub(1),lb(2):ub(2)))       ! Computed solutions
  ALLOCATE(resids(lb(1):ub(1),lb(2):ub(2)))  ! Residuals
  ALLOCATE(errs(lb(1):ub(1),lb(2):ub(2)))    ! Errors
  ALLOCATE(resid_it(0:nits))
  ALLOCATE(err_it(0:nits))
!
  f(s(1):e(1),s(2):e(2)) = rhs(xgrid(s(1):e(1)),ygrid(s(2):e(2)))
  v(s(1):e(1),s(2):e(2)) = exact(xgrid(s(1):e(1)),ygrid(s(2):e(2)))
  CALL exchange(comm_cart, f)
  CALL exchange(comm_cart, v)
!
!   Residuals of exact solutions
  resids = get_resids(mat,v,f)
  resid = norm2_vec(resids, comm_cart)
!
!   Jacobi iteration loop
!
  IF(me.EQ.0) WRITE(*,'(/a6,t14,a,t34,a)') 'it', 'residual norm', 'discretization error'
  u = 0.0d0
  CALL exchange(comm_cart, u)
  resids = get_resids(mat,u,f)
  errs  = u-v
  resid_it(0) = norm2_vec(resids, comm_cart)
  err_it(0)   = norm2_vec(errs, comm_cart)
  it_skip = MAX(1,nits/10)
!
  CALL timera(0, 'Jacobi')
  DO it=1,nits
     CALL jacobi(mat, omega, 1, u, f)
     CALL exchange(comm_cart, u)
     resids = get_resids(mat,u,f)
     errs  = u-v
     resid_it(it) = norm2_vec(resids, comm_cart)
     err_it(it)   = norm2_vec(errs, comm_cart)
     IF(me.EQ.0 .AND. MODULO(it,it_skip).EQ.0 ) THEN 
        WRITE(*,'(i6,4(1pe12.3))') it, resid_it(it), resid_it(it)/resid_it(it-1),&
          &    err_it(it), err_it(it)/err_it(it-1)
     END IF
  END DO
  CALL timera(1, 'Jacobi')
!================================================================================
!                        9.0  Epilogue
  CALL h5file
!
  CALL timera(9, '')
  CALL MPI_FINALIZE(ierr)
CONTAINS
  SUBROUTINE disp(str, arr)
    CHARACTER(len=*), INTENT(in) :: str
    REAL(rkind), INTENT(in)      :: arr(:,:)
    INTEGER :: j
    WRITE(*,'(/a)') str
    DO j=1,SIZE(arr,2)
       WRITE(*,'(10f8.3)') arr(:,j)
    END DO
  END SUBROUTINE disp
!
  SUBROUTINE h5file
    USE futils
    CHARACTER(len=128) :: file='test_jacobi.h5'
    INTEGER :: fid
    CALL creatf(file, fid, real_prec='d', mpicomm=comm_cart)
    CALL putarr(fid, '/xgrid', xgrid, ionode=0)  ! only rank 0 does IO
    CALL putarr(fid, '/ygrid', ygrid, ionode=0)  ! only rank 0 does IO
!
    CALL putarrnd(fid, '/f',      f,     (/1,2/), garea=(/1,1/))
    CALL putarrnd(fid, '/v',      v,     (/1,2/), garea=(/1,1/))
    CALL putarrnd(fid, '/u',      u,     (/1,2/), garea=(/1,1/))
    CALL putarrnd(fid, '/errs',   errs,  (/1,2/), garea=(/1,1/))
    CALL putarrnd(fid, '/resids', resids,(/1,2/), garea=(/1,1/))
!
    CALL putarr(fid, '/resid', resid_it, ionode=0)
    CALL putarr(fid, '/error', err_it, ionode=0)
!
    CALL putmat(fid, '/MAT', mat)
    CALL closef(fid)
  END SUBROUTINE h5file
!
  FUNCTION outerprod(x, y) RESULT(r)
!
!  outer product
!
    REAL(rkind), INTENT(in) :: x(:), y(:)
    REAL(rkind)             :: r(SIZE(x),SIZE(y))
    INTEGER :: i, j
    DO j=1,SIZE(y)
       DO i=1,SIZE(x)
          r(i,j) = x(i)*y(j)
       END DO
    END DO
  END FUNCTION outerprod
!
  FUNCTION rhs(x,y)
    REAL(rkind), INTENT(in) :: x(:), y(:)
    REAL(rkind)             :: rhs(SIZE(x),SIZE(y))
    rhs = -10.d0*pi**2 * outerprod(SIN(pi*x), SIN(3.d0*pi*y))
  END FUNCTION rhs
!
  FUNCTION exact(x,y)
    REAL(rkind), INTENT(in) :: x(:), y(:)
    REAL(rkind)             :: exact(SIZE(x),SIZE(y))
    exact = outerprod(SIN(pi*x), SIN(3.d0*pi*y))
  END FUNCTION exact
END PROGRAM main