!>
!> @file test_intergrid1.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>
!>
PROGRAM main
!
!   Test implementation of (parallel) matrix-free
!
  USE iso_fortran_env, ONLY : rkind => real64
  USE parmg, ONLY    : grid2_type, init_restrict, coarse, get_lmax, &
       &               exchange, prolong, restrict, disp, norm_vec
  USE pputils2, ONLY : dist1d
  USE gvector, ONLY  : gvector_2d,OPERATOR(-)
  USE futils
  USE mpi
  IMPLICIT NONE
!
  INTEGER, PARAMETER :: ndims=2
  INTEGER                   :: ierr, me, npes
  INTEGER, DIMENSION(ndims) :: dims=[0,0]
  INTEGER, DIMENSION(ndims) :: lmax, coords, comm1d
  LOGICAL, DIMENSION(ndims) :: periods=[.FALSE.,.FALSE.]
  LOGICAL                   :: reorder =.FALSE.
  INTEGER                   :: comm_cart, comm_futils
!
  INTEGER           :: fin
  CHARACTER(len=64) :: filein = 'test_intergrid0.h5'
  CHARACTER(len=64) :: dsname
  CHARACTER(len=4)  :: prb
  LOGICAL           :: nldebug
!
  INTEGER                       :: nx, ny, levels
  TYPE(grid2_type), ALLOCATABLE :: grids(:), new_grids(:)
  INTEGER, DIMENSION(ndims)     :: e, s, npt_glob, npt_loc, npt_loc_min
!
  CHARACTER(len=64) :: str
  REAL(rkind)       :: err
  INTEGER           :: i, k, l
!--------------------------------------------------------------------------------
!                              1.0 Prologue
!
!   Init MPI and setup 2D grid topology
!
  CALL mpi_init(ierr)
  CALL mpi_comm_size(MPI_COMM_WORLD, npes, ierr)
  CALL mpi_comm_rank(MPI_COMM_WORLD, me, 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_cart_coords(comm_cart, me, ndims, coords, ierr)
  CALL mpi_cart_sub(comm_cart, [.TRUE.,.FALSE.], comm1d(1), ierr)
  CALL mpi_cart_sub(comm_cart, [.FALSE.,.TRUE.], comm1d(2), ierr)
!
  IF( me .EQ. 0 ) WRITE(*,'(a,i3,i3)') '2d processor grid', dims
!
!   Get nx, ny, levels from h5 file created by test_intergrid0
!
  IF( command_argument_count() > 0 ) THEN
     CALL get_command_argument(1, filein)
  END IF
  IF(me.EQ.0) WRITE(*,'(a,a)') 'filein = ', TRIM(filein)
!
  CALL mpi_comm_dup(comm_cart, comm_futils, ierr)
  CALL openf(filein, fin, mpicomm=comm_futils)
  CALL getatt(fin, '/', 'NX', nx, ierr)
  CALL getatt(fin, '/', 'NY', ny, ierr)
  CALL getatt(fin, '/', 'LEVELS', levels, ierr)
  CALL getatt(fin, '/', 'PRB', prb, ierr)
  CALL getatt(fin, '/', 'NLDEBUG', nldebug, ierr)
  IF(me.EQ.0) WRITE(*,'(a,a,3i5,l3)') 'prb, nx, ny, levels: ', prb, nx, ny, &
       &                              levels, nldebug
!--------------------------------------------------------------------------------
!                              2.0 Read (f,v) from h5 file
!
  ALLOCATE(grids(levels))
!
!    Partition on finest grid
!
  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
     lmax(i) = get_lmax(s(i), npt_loc(i), 1, comm1d(i))
  END DO
  npt_loc = e-s+1
  IF(me.EQ.0) WRITE(*,'(a,2i4)') 'lmax', lmax
!
!    Partition on coaser grids
!
  DO l=1,levels
     IF(l.GT.1) THEN
        CALL coarse(s,e)
        npt_loc = e-s+1
        CALL mpi_allreduce(npt_loc, npt_loc_min, 2, MPI_INTEGER, &
             &             MPI_MIN, comm_cart, ierr)
        CALL mpi_allreduce(e, npt_glob, 2, MPI_INTEGER, MPI_MAX, &
             &             comm_cart, ierr)
        npt_glob = npt_glob+1
     END IF
     WRITE(str,'(a,i3,a)') 'Partition at level', l, ': start. index ='
     CALL disp(TRIM(str), s, comm_cart)
     IF(me.EQ.0) THEN 
        WRITE(*,'(a,2i6)') 'npt_glob   ', npt_glob
        WRITE(*,'(a,2i6)') 'npt_loc_min', npt_loc_min
     END IF
     grids(l)%s = s
     grids(l)%e = e
     grids(l)%npt = npt_glob
     grids(l)%f = gvector_2d(s, e, [1,1])
     grids(l)%v = gvector_2d(s, e, [1,1])
     ALLOCATE(grids(l)%x(0:npt_glob(1)-1))  ! Global coords (x,y)
     ALLOCATE(grids(l)%y(0:npt_glob(2)-1))
     WRITE(dsname,'("/mglevels/level.",i2.2)') l
     CALL getarr(fin, TRIM(dsname)//"/x", grids(l)%x)
     CALL getarr(fin, TRIM(dsname)//"/y", grids(l)%y)
     CALL getarrnd(fin, TRIM(dsname)//"/v", grids(l)%v%val, [1,2], garea=[1,1])
     CALL getarrnd(fin, TRIM(dsname)//"/f", grids(l)%f%val, [1,2], garea=[1,1])
  END DO
!--------------------------------------------------------------------------------
!                              3.0 Parallel intergrid transfer
!
  ALLOCATE(new_grids(levels))
  CALL copy_grids(grids, new_grids)
!
!   Set up restriction stencil
!
  DO l=2,levels
     CALL init_restrict(new_grids(l), prb, comm_cart)
  END DO
!
!   Prolongation of v
!
  DO l=levels-1,1,-1
     CALL exchange(comm_cart, grids(l+1)%v)
     CALL prolong(grids(l+1)%v, new_grids(l)%v)
     IF(nldebug) THEN
        IF(me.EQ.0) WRITE(*,'(a)') '====='
        DO k=0,npes
           IF(me.EQ.k) THEN
              s = grids(l+1)%f%s
              e = grids(l+1)%f%e
              WRITE(*,'(a,i2)') 'reference vbar on proc.', me
              DO i=s(1),e(1)
                 WRITE(*,'(10f8.3)') grids(l+1)%v%val(i,s(2):e(2))
              END DO
              s = grids(l)%f%s
              e = grids(l)%f%e
              WRITE(*,'(a,i2)') 'reference v on proc.', me
              DO i=s(1),e(1)
                 WRITE(*,'(10f8.3)') grids(l)%v%val(i,s(2):e(2))
              END DO
              WRITE(*,'(a,i2)') 'compute v on proc.', me
              DO i=s(1),e(1)
                 WRITE(*,'(10f8.3)') new_grids(l)%v%val(i,s(2):e(2))
              END DO
           END IF
           CALL mpi_barrier(comm_cart, ierr)
        END DO
     END IF
     err = norm_vec(new_grids(l)%v-grids(l)%v, comm_cart, 0)
     IF(me.EQ.0) WRITE(*,'(a,i3,1pe12.3)') 'Error of prolongation: ', l, err
  END DO
!
!   Restriction of f
!
  DO l=2,levels
     CALL exchange(comm_cart, grids(l-1)%f)
     CALL restrict(new_grids(l)%restrict_mat, grids(l-1)%f, new_grids(l)%f)
     IF(nldebug) THEN
        IF(me.EQ.0)  WRITE(*,'(a)') '====='
        DO k=0,npes
           IF(me.EQ.k) THEN
              s = grids(l-1)%f%s
              e = grids(l-1)%f%e
              WRITE(*,'(a,i2)') 'reference f on proc.', me
              DO i=s(1),e(1)
                 WRITE(*,'(10f8.3)') grids(l-1)%f%val(i,s(2):e(2))
              END DO
              s = grids(l)%f%s
              e = grids(l)%f%e
              WRITE(*,'(a,i2)') 'reference fbar on proc.', me
              DO i=s(1),e(1)
                 WRITE(*,'(10f8.3)') grids(l)%f%val(i,s(2):e(2))
              END DO
              WRITE(*,'(a,i2)') 'compute fbar on proc.', me
              DO i=s(1),e(1)
                 WRITE(*,'(10f8.3)') new_grids(l)%f%val(i,s(2):e(2))
              END DO
           END IF
           CALL mpi_barrier(comm_cart, ierr)
        END DO
     END IF
     err = norm_vec(new_grids(l)%f-grids(l)%f, comm_cart, 0)
     IF(me.EQ.0) WRITE(*,'(a,i3,1pe12.3)') 'Error of restriction: ', l, err
  END DO
!--------------------------------------------------------------------------------
!                              9.0 Epilogue
!
  CALL closef(fin)  
  CALL mpi_finalize(ierr)
!
CONTAINS
  SUBROUTINE copy_grids(g1, g2)
    TYPE(grid2_type), INTENT(in)  :: g1(:)
    TYPE(grid2_type), INTENT(out) :: g2(:)
    INTEGER :: l
    DO l=1,SIZE(g1)
       g2(l)%s = g1(l)%s
       g2(l)%e = g1(l)%e
       g2(l)%npt_loc = g1(l)%npt_loc
       g2(l)%npt = g1(l)%npt
       ALLOCATE(g2(l)%x(0:g2(l)%npt(1)-1)); g2(l)%x = g1(l)%x
       ALLOCATE(g2(l)%y(0:g2(l)%npt(2)-1)); g2(l)%y = g1(l)%y
       g2(l)%v = gvector_2d(g1(l)%v%s, g1(l)%v%e, g1(l)%v%g); g2(l)%v%val = g1(l)%f%val
       g2(l)%f = gvector_2d(g1(l)%f%s, g1(l)%f%e, g1(l)%f%g); g2(l)%f%val = g1(l)%f%val
    END DO
  END SUBROUTINE copy_grids
END PROGRAM main