!>
!> @file test_jacobig.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/>.
!>
!> @authors
!> (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 fdmat_mod, ONLY : fdmat, ibc_fdmat, ibc_rhs
  USE pputils2, ONLY  : dist1d, timera, hostlist
  USE gvector, ONLY   : gvector_2d, ASSIGNMENT(=), OPERATOR(-)
  USE parmg, ONLY     : grid2_type, create_grid, jacobi, exchange, get_resids, norm_vec
  USE stencil, ONLY   : stencil_2d, 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, DIMENSION(ndims) :: e0, s0, e, s, npt_glob, npt_loc
!
  REAL(rkind), ALLOCATABLE  :: xgrid(:), ygrid(:)
  INTEGER, ALLOCATABLE      :: id(:,:)
  REAL(rkind)               :: dx, dy
  INTEGER                   :: npoints ! Number of points in FD stencil
!
  TYPE(gvector_2d)          :: v_exact, resids, errs
  REAL(rkind), ALLOCATABLE  :: resid_it(:), err_it(:)
  INTEGER, DIMENSION(ndims) :: g
  INTEGER                   :: i, it, it_skip
!
  INTEGER :: levels=1
  TYPE(grid2_type), ALLOCATABLE :: grids(:)
!
!   Input quantities
!
  CHARACTER(len=4) :: prb='dddd'
  INTEGER          :: nx=4, ny=4    ! Number of intervals
  INTEGER          :: kx=1, ky=1
  REAL(rkind)      :: Lx=1.0, Ly=1.0
  REAL(rkind)      :: icrosst=1.0, beta=0.0, miome=200.0
  REAL(rkind)      :: omega=1.0d0
  INTEGER          :: nits=100, nu=1
!
  NAMELIST /in/ prb, nx, ny, kx, ky, Lx, Ly, icrosst, beta, &
       &        miome, omega, nits, nu
!================================================================================
!                        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(prb, LEN(prb), MPI_CHARACTER, 0, comm_cart, ierr)
  CALL mpi_bcast(kx, 1, MPI_INTEGER, 0, comm_cart, ierr)
  call mpi_bcast(ky, 1, MPI_INTEGER, 0, comm_cart, ierr)
  CALL mpi_bcast(icrosst, 1, MPI_DOUBLE_PRECISION, 0, comm_cart, ierr)
  CALL mpi_bcast(Lx,1,MPI_DOUBLE_PRECISION,0,comm_cart, ierr)
  CALL mpi_bcast(Ly,1,MPI_DOUBLE_PRECISION,0,comm_cart, ierr)
  CALL mpi_bcast(beta,1,MPI_DOUBLE_PRECISION,0,comm_cart, ierr)
  CALL mpi_bcast(miome, 1, MPI_DOUBLE_PRECISION, 0, comm_cart, ierr)
  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(nu, 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 = Lx/REAL(nx)
  dy = Ly/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')
!
!   Create grid structure
!
  ALLOCATE(grids(levels))
  npoints = 9             ! Size of FD stencil
  ALLOCATE(id(npoints,2))
  id=RESHAPE([ 0, -1,  0,  1, -1, 1, -1, 0, 1,  & 
               0, -1, -1, -1,  0, 0,  1, 1, 1], &
              [npoints,2])
  CALL create_grid(xgrid, ygrid, s, e, id, prb, grids, comm_cart)
!================================================================================
!                        3.0  FD Operator
!
  CALL timera(0, 'Laplacian')
!
  CALL fdmat(grids(1),  fdense, icrosst, grids(1)%fdmat)
  CALL ibc_fdmat(grids(1)%fdmat, prb)
!
  CALL timera(1, 'Laplacian')
!================================================================================
!                        4.0 RHS and exact solution
!
!   Allocate memory
!
  s0 = grids(1)%s0; e0 = grids(1)%e0
  s  = grids(1)%s;  e  = grids(1)%e
  g = [1,1]
  v_exact = gvector_2d(s, e, g) ! Exact solutions   
  errs    = gvector_2d(s, e, g) ! Disc. errors
  resids  = gvector_2d(s, e, g) ! Residuals 
  ALLOCATE(resid_it(0:nits))
  ALLOCATE(err_it(0:nits))
!
!   Set RHS at the finest grid. Impose Dirichlet BC.
!
  grids(1)%f = frhs(xgrid(s(1):e(1)),ygrid(s(2):e(2)))
  CALL ibc_rhs(grids(1)%f, s0, e0, prb)
!
!   Exact solutions
!
  v_exact = fexact(xgrid(s(1):e(1)),ygrid(s(2):e(2)))
!================================================================================
!                        5.0  Jacobi iteration loop
!
  IF(me.EQ.0) WRITE(*,'(/a6,t14,a,t34,a)') 'it', 'residual norm', 'discretization error'
  grids(1)%v = 0.0d0
  resids = get_resids(comm_cart, grids(1)%fdmat, grids(1)%v, grids(1)%f)
  errs  = grids(1)%v - v_exact
  resid_it(0) = norm_vec(resids, comm_cart, root=0)
  err_it(0)   = norm_vec(errs, comm_cart, root=0)
  it_skip = MAX(1,nits/10)
!
  CALL timera(0, 'Jacobi')
  DO it=1,nits
     CALL jacobi(comm_cart, grids(1)%fdmat, omega, nu, grids(1)%v, grids(1)%f)
     resids = get_resids(comm_cart, grids(1)%fdmat, grids(1)%v, grids(1)%f)
     errs  = grids(1)%v - v_exact
     resid_it(it) = norm_vec(resids, comm_cart, root=0)
     err_it(it)   = norm_vec(errs, comm_cart, root=0)
     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
!
!+++
  FUNCTION fdense(x)
!
!   Return density
!
    REAL(rkind), INTENT(in) :: x(:)
    REAL(rkind)              :: fdense(SIZE(x))
    fdense(:) = 0.5d0*beta*miome*EXP( -(x(:)-Lx/3d0)**2d0/(Lx/2d0)**2d0 );
  END FUNCTION fdense
!+++
  FUNCTION fexact(x,y)
!
!   Return analytical solution
!
    REAL(rkind), INTENT(in) :: x(:), y(:)
    REAL(rkind) :: fexact(SIZE(x),SIZE(y))
    REAL(rkind) :: c
    INTEGER :: j
    IF(prb.EQ.'dddd') THEN
       DO j=1,SIZE(y)
          c = SIN(2.0d0*pi*ky*y(j)/Ly)
          fexact(:,j) = SIN(2.0d0*pi*kx*x(:)/Lx)*c
       END DO
    ELSE IF (prb.EQ.'nndd') THEN
       DO j=1,SIZE(y)
          c = SIN(2.0d0*pi*ky*y(j)/Ly)
          fexact(:,j) = COS(2.0d0*pi*kx*x(:)/Lx)*c
       END DO
    END IF
  END FUNCTION fexact
!+++
  SUBROUTINE h5file
    USE futils
    CHARACTER(len=128) :: file='test_jacobig.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', grids(1)%f%val, (/1,2/), garea=g)
    CALL putarrnd(fid, '/v', v_exact%val, (/1,2/), garea=g)
    CALL putarrnd(fid, '/u', grids(1)%v%val, (/1,2/), garea=g)
    CALL putarrnd(fid, '/errs',   errs%val,  (/1,2/), garea=(/1,1/))
    CALL putarrnd(fid, '/resids', resids%val,(/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', grids(1)%fdmat)
    CALL closef(fid)
  END SUBROUTINE h5file
!+++
  FUNCTION frhs(x,y)
!
!   Return RHS
!
    REAL(rkind), INTENT(in) :: x(:), y(:)
    REAL(rkind) :: frhs(SIZE(x),SIZE(y))
    REAL(rkind) :: c, s, d(SIZE(x))
    REAL(rkind) :: corr
    INTEGER :: j
    corr = 1.d0+icrosst**2/4.0d0
    d(:) = fdense(x(:))
    IF(prb.EQ.'dddd') THEN
       DO j=1,SIZE(y)
          c = COS(2.0d0*pi*ky*y(j)/Ly)
          s = SIN(2.0d0*pi*ky*y(j)/Ly)
          frhs(:,j)=-4.0d0*pi*pi*( (kx*kx/Lx/Lx+corr*ky*ky/Ly/Ly) * SIN(2.0d0*pi*kx*x(:)/Lx)*s &
               &     -icrosst*kx*ky/Lx/Ly * COS(2.0d0*pi*kx*x(:)/Lx)*c ) &
               &     + d(:) * SIN(2.0d0*pi*kx*x(:)/Lx)*s
       END DO
    ELSE IF (prb.EQ.'nndd') THEN
       DO j=1,SIZE(y)
          c = COS(2.0d0*pi*ky*y(j)/Ly)
          s = SIN(2.0d0*pi*ky*y(j)/Ly)
          frhs(:,j)=-4.0d0*pi*pi*( (kx*kx/Lx/Lx+corr*ky*ky/Ly/Ly) * COS(2.0d0*pi*kx*x(:)/Lx)*s &
               &     +icrosst*kx*ky/Lx/Ly * SIN(2.0d0*pi*kx*x(:)/Lx)*c ) &
               &     + d(:) * COS(2.0d0*pi*kx*x(:)/Lx)*s
       END DO
    END IF
  END FUNCTION frhs
!+++!
  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