!>
!> @file test_mg2d.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>
!>
PROGRAM main
!
!   Test 2d multigrid: Cartesian case
!
  USE multigrid
  USE csr
  IMPLICIT NONE
  INCLUDE 'mpif.h'
!
  DOUBLE PRECISION, PARAMETER :: PI=4.0d0*ATAN(1.0d0)
  INTEGER, DIMENSION(2) :: n, nidbas, ngauss, alpha
  DOUBLE PRECISION :: kx=4.d0, ky=3.d0, sigma=10.0d0
  CHARACTER(len=4) :: prb='poly'
  INTEGER :: levels=1, nu1=1, nu2=1, mu=1, nu0=1, nits
  CHARACTER(len=4) :: relax='jac '
  DOUBLE PRECISION :: omega=2.0d0/3.0d0, tol
  LOGICAL :: nlfixed=.FALSE.
  DOUBLE PRECISION :: t0, tsetup(2), tdirect, tbsolve, titer, titer_per_step
  DOUBLE PRECISION :: resid_direct, errdisc_direct
  DOUBLE PRECISION :: norma, normb
!
  DOUBLE PRECISION, ALLOCATABLE :: x(:), y(:)
  DOUBLE PRECISION :: dx, dy
  INTEGER :: ix, iy
!
  DOUBLE PRECISION, ALLOCATABLE :: sol_anal_grid(:,:)
  DOUBLE PRECISION, ALLOCATABLE :: sol_calc_grid(:,:)
  DOUBLE PRECISION, ALLOCATABLE :: rresid(:), resid(:), errdisc(:)
!
  DOUBLE PRECISION, ALLOCATABLE, TARGET :: sol_direct(:,:)
  DOUBLE PRECISION, POINTER     :: sol_direct_1d(:)
  DOUBLE PRECISION, ALLOCATABLE :: sol_direct_grid(:,:)
!
  INTEGER :: ierr, me
  INTEGER :: l, nterms
  INTEGER :: its
!
  TYPE(grid2d), ALLOCATABLE :: grids(:)
  TYPE(mg_info)             :: info
!
  NAMELIST /newrun/ n, nidbas, ngauss, kx, ky, sigma, alpha, levels, prb, &
       &            nu1, nu2, mu, nu0, relax, nits, tol, nlfixed, omega
!--------------------------------------------------------------------------------
!                    1.   Prologue
!
  CALL mpi_init(ierr)
  CALL mpi_comm_rank(MPI_COMM_WORLD, me, ierr)
!
!   Inputs
!
  n = (/8, 8/)
  nidbas=(/3,3/)
  ngauss=(/2,2/)
  alpha = (/0,0/)
  kx=4
  ky=3
  sigma=10.0d0
  levels=2
  prb='poly'
  nu1 = 1
  nu2 = 1
  mu = 1
  nu0 = 1
  nits = 10
  tol = 1.e-8
  relax = 'jac'
  nlfixed= .FALSE.
!
  IF(me.EQ.0) THEN
     READ(*,newrun)
     WRITE(*,newrun)
  END IF
     CALL mpi_bcast(n, 2, MPI_INTEGER, 0, MPI_COMM_WORLD, ierr)
     CALL mpi_bcast(nidbas, 2, MPI_INTEGER, 0, MPI_COMM_WORLD, ierr)
     CALL mpi_bcast(ngauss, 2, MPI_INTEGER, 0, MPI_COMM_WORLD, ierr)
     CALL mpi_bcast(alpha, 2, MPI_DOUBLE_PRECISION, 0, MPI_COMM_WORLD, ierr)
     CALL mpi_bcast(kx, 1, MPI_DOUBLE_PRECISION, 0, MPI_COMM_WORLD, ierr)
     CALL mpi_bcast(ky, 1, MPI_DOUBLE_PRECISION, 0, MPI_COMM_WORLD, ierr)
     CALL mpi_bcast(sigma, 1, MPI_DOUBLE_PRECISION, 0, MPI_COMM_WORLD, ierr)
     CALL mpi_bcast(levels, 1, MPI_INTEGER, 0, MPI_COMM_WORLD, ierr)
     CALL mpi_bcast(prb, LEN(prb), MPI_CHARACTER, 0, MPI_COMM_WORLD, ierr)
     CALL mpi_bcast(nu1, 1, MPI_INTEGER, 0, MPI_COMM_WORLD, ierr)
     CALL mpi_bcast(nu2, 1, MPI_INTEGER, 0, MPI_COMM_WORLD, ierr)
     CALL mpi_bcast(nu0, 1, MPI_INTEGER, 0, MPI_COMM_WORLD, ierr)
     CALL mpi_bcast(mu, 1, MPI_INTEGER, 0, MPI_COMM_WORLD, ierr)
     CALL mpi_bcast(nits, 1, MPI_INTEGER, 0, MPI_COMM_WORLD, ierr)
     CALL mpi_bcast(relax, 4, MPI_CHARACTER, 0, MPI_COMM_WORLD, ierr)
     CALL mpi_bcast(nlfixed, 1, MPI_LOGICAL, 0, MPI_COMM_WORLD, ierr)
     CALL mpi_bcast(omega, 1, MPI_DOUBLE_PRECISION, 0, MPI_COMM_WORLD, ierr)
     CALL mpi_bcast(tol, 1, MPI_DOUBLE_PRECISION, 0, MPI_COMM_WORLD, ierr)
!
!   Adjust number of levels and fill mg info.
!
  levels = MIN(levels, get_lmax(n(1)), get_lmax(n(2)))
  info%nu1 = nu1
  info%nu2 = nu2
  info%mu = mu
  info%nu0 = nu0
  info%levels = levels
  info%relax = relax
  info%omega = omega
!
!   Create grids
!
  t0 = mpi_wtime()
!
  dx = 1.0d0/REAL(n(1),8)
  dy = 1.0d0/REAL(n(2),8)
  ALLOCATE(x(0:n(1)), y(0:n(2)))
  DO ix=0,n(1)
     x(ix) = ix*dx
  END DO
  DO iy=0,n(2)
     y(iy) = iy*dy
  END DO
!
  ALLOCATE(grids(levels))
  CALL create_grid(x, y, nidbas, ngauss, alpha, grids)
  WRITE(*,'(5a6)') 'l', 'nx', 'ny', 'rx', 'ry'
  WRITE(*,'(5i6)') (l, grids(l)%n, grids(l)%rank, l=1,levels)
!
!   Construct RHS and set BC only on the finest grid
!
  CALL disrhs(grids(1)%spl, grids(1)%f, rhs)
  CALL ibcrhs(grids(1), grids(1)%f)
!
!   Build FE matrices and set BC
!
  nterms = 3
  DO l=1,levels
     CALL femat(grids(l)%spl, grids(l)%mata, coefeq, nterms)
     CALL ibcmat(grids(l), grids(l)%mata)
     CALL to_mat(grids(l)%mata)
  END DO
!
!   Set BC on grid transfer matrices
!
  CALL ibc_transf(grids,1,3)
  CALL ibc_transf(grids,2,3)
!
  tsetup(1) = mpi_wtime()-t0
!
!   Clear and rebuild FE matrices and set BC
!
  t0 = mpi_wtime()
  nterms = 2
  DO l=1,levels
     CALL clear_mat(grids(l)%mata)
     CALL femat(grids(l)%spl, grids(l)%mata, coefeq, nterms, noinit=.TRUE.)
     CALL ibcmat(grids(l), grids(l)%mata)
  END DO
  tsetup(2) = mpi_wtime()-t0
!--------------------------------------------------------------------------------
!                   1.  Analytical solution (at the finest grid, l=1)
!
  ALLOCATE(sol_anal_grid(0:n(1),0:n(2)))
  sol_anal_grid = sol(grids(1)%x, grids(1)%y)
!--------------------------------------------------------------------------------
!                   2.  Direct solution (at the finest grid, l=1)
!
  WRITE(*,'(//a)') 'Direct solution for the finest grid problem'
!
  ALLOCATE(sol_direct(SIZE(grids(1)%v,1), SIZE(grids(1)%v,2)), &
       &         source=grids(1)%f)
  ALLOCATE(sol_direct_grid(0:n(1),0:n(2)))
!
  sol_direct_1d(1:SIZE(grids(1)%v1d)) => sol_direct
!
!
  PRINT*, 'shape of sol_direct', SHAPE(sol_direct)
  PRINT*, 'shape of sol_direct_1d', SHAPE(sol_direct_1d)
!
  t0 = mpi_wtime()
  sol_direct = grids(1)%f
  CALL direct_solve(grids(1), sol_direct_1d, debug=.FALSE.)
  tdirect = mpi_wtime()-t0
!
  t0 = mpi_wtime()
  sol_direct = grids(1)%f
  CALL direct_solve(grids(1), sol_direct_1d, debug=.FALSE.)
  resid_direct = residue(grids(1)%mata, grids(1)%f1d, sol_direct_1d)
  errdisc_direct = disc_err(grids(1)%spl, sol_direct, sol)
!
  tbsolve = mpi_wtime()-t0
!
  CALL gridval(grids(1)%spl, grids(1)%x, grids(1)%y, sol_direct_grid, &
       &       [0,0], sol_direct)
!
  WRITE(*,'(a,2(1pe12.3))') 'Discretization error and residue =', &
       &        errdisc_direct, resid_direct
  WRITE(*,'(a,2(1pe12.3))') 'Relative discretization errors', &
       &    NORM2(sol_anal_grid-sol_direct_grid) / NORM2(sol_anal_grid)
!
  WRITE(*, '(a,1pe12.3)') 'Frobenius norm of A', matnorm(grids(1)%mata)
  WRITE(*, '(a,1pe12.3)') 'Infinity norm of A ', matnorm(grids(1)%mata, 'inf')
  WRITE(*, '(a,1pe12.3)') '1 norm of A        ', matnorm(grids(1)%mata, '1')
!--------------------------------------------------------------------------------
!                   3.  Test multigrid V-cycle
!
  WRITE(*,'(/a)') 'Multigrid MG V-cycles ...'
  ALLOCATE(sol_calc_grid(0:n(1),0:n(2)))
  ALLOCATE(errdisc(0:nits))
  ALLOCATE(resid(0:nits))
  ALLOCATE(rresid(0:nits))
!
!   Norm of A and b
!
  norma = matnorm(grids(1)%mata)
  normb = NORM2(grids(1)%f1d)
!
!   Initial guess
!
  IF(nlfixed) THEN
     grids(1)%v = sol_direct
  ELSE
     grids(1)%v = 0.0d0
  END IF
!
  errdisc(0) = disc_err(grids(1)%spl, grids(1)%v, sol)
  resid(0) = residue(grids(1)%mata, grids(1)%f1d, grids(1)%v1d)
  rresid(0) = resid(0)  /  ( norma*NORM2(grids(1)%v1d) + normb )
  WRITE(*,'(a4,3(a12,a8))') 'its', 'residue', 'ratio', 'disc. err', 'ratio', &
       &                           'rel. resid', 'ratio'
  WRITE(*,'(i4,3(1pe12.3,8x))') 0, resid(0), errdisc(0), rresid(0)
!
!   Iterations
!
  t0 = mpi_wtime()
  DO its=1,nits
     CALL mg(grids, info, 1)
     errdisc(its) = disc_err(grids(1)%spl, grids(1)%v, sol)
     resid(its) = residue(grids(1)%mata, grids(1)%f1d, grids(1)%v1d)
     rresid(its) = resid(its)  / ( norma*NORM2(grids(1)%v1d) + normb )
     WRITE(*,'((i4,3(1pe12.3,0pf8.2)))')  its, &
          &   resid(its),   resid(its)/resid(its-1), &
          &   errdisc(its), errdisc(its)/errdisc(its-1), &
          &   rresid(its),  rresid(its)/rresid(its-1)
     IF(resid(its) .LE. tol) EXIT
  END DO
  nits = MIN(nits,its)
  titer = mpi_wtime() - t0
  titer_per_step = titer/REAL(its,8)
!
  CALL gridval(grids(1)%spl, grids(1)%x, grids(1)%y, sol_calc_grid, &
       &       [0,0], grids(1)%v)
!--------------------------------------------------------------------------------
!                   9.  Epilogue
!
!   Display timing
!
  WRITE(*,'(a,2(1pe12.3))') 'Set up time (s)            ', tsetup
  WRITE(*,'(a,2(1pe12.3))') 'Direct and solve time (s)  ', tdirect, tbsolve
  WRITE(*,'(a,1pe12.3,i5)') 'Iter time (s)              ', titer, nits
!
!   Creata HDF5 file
!
  IF(me.EQ.0) CALL h5file
!
  CALL mpi_finalize(ierr)
!--------------------------------------------------------------------------------
CONTAINS
!+++
  FUNCTION rhs(x, y)
!
! Return problem RHS
!
    DOUBLE PRECISION, INTENT(in) :: x, y
    DOUBLE PRECISION :: rhs
    DOUBLE PRECISION :: x2, y2
    SELECT CASE(TRIM(prb))
    CASE('poly')
       x2 = x*x; y2 = y*y;
       rhs = 2.d0 * ( (1.0d0-6.d0*x2)*y2*(1.d0-y2) + &
            &         (1.0d0-6.d0*y2)*x2*(1.d0-x2)  )
    CASE('trig')
       rhs = SIN(PI*kx*x)*SIN(PI*ky*y)
    END SELECT
  END FUNCTION rhs
!+++
  FUNCTION sol(x, y)
!
! Return exact problem solution
!
    DOUBLE PRECISION, INTENT(in) :: x(:), y(:)
    DOUBLE PRECISION :: sol(SIZE(x),SIZE(y))
    DOUBLE PRECISION :: c
    DOUBLE PRECISION  :: x2(SIZE(x)), y2(SIZE(y))
    INTEGER :: j
    SELECT CASE(TRIM(prb))
    CASE('poly')
       x2 = x*x; y2 = y*y;
       DO j=1,SIZE(y)
          c = y2(j)*(y2(j)-1.d0)
          sol(:,j) = c*x2(:)*(1.0d0-x2(:))
       END DO
    CASE('trig')
       DO j=1,SIZE(y)
          c = SIN(PI*ky*y(j)) / (PI**2*(kx**2+ky**2) + sigma**2)
          sol(:,j) = c * SIN(PI*kx*x(:))
       END DO
    END SELECT
  END FUNCTION sol
!+++
  SUBROUTINE coefeq(x, y, idt, idw, c)
!
! Weak form = Int( \nabla(w).\nabla(t) + \sigma.t.w) dV)
!
    DOUBLE PRECISION, INTENT(in)  :: x, y
    INTEGER, INTENT(out)           :: idt(:,:), idw(:,:)
    DOUBLE PRECISION, INTENT(out) :: c(:)
!
    c(1) = 1.0d0
    idt(1,1) = 1
    idt(1,2) = 0
    idw(1,1) = 1
    idw(1,2) = 0
!
    c(2) = 1.0d0
    idt(2,1) = 0
    idt(2,2) = 1
    idw(2,1) = 0
    idw(2,2) = 1
!
    c(3) = sigma
    idt(3,1) = 0
    idt(3,2) = 0
    idw(3,1) = 0
    idw(3,2) = 0

  END SUBROUTINE coefeq
!+++
  SUBROUTINE h5file
    USE futils
    CHARACTER(len=128) :: file='test_mg2d.h5'
    INTEGER :: fid
    INTEGER :: l
    CHARACTER(len=64) :: dsname
    CALL creatf(file, fid, real_prec='d')
    CALL attach(fid, '/', 'NX', n(1))
    CALL attach(fid, '/', 'NY', n(2))
    CALL attach(fid, '/', 'NIDBAS1', nidbas(1))
    CALL attach(fid, '/', 'NIDBAS2', nidbas(2))
    CALL attach(fid, '/', 'KX', kx)
    CALL attach(fid, '/', 'KY', ky)
    CALL attach(fid, '/', 'SIGMA', sigma)
    CALL attach(fid, '/', 'ALPHA1', alpha(1))
    CALL attach(fid, '/', 'ALPHA2', alpha(2))
    CALL attach(fid, '/', 'RELAX', relax)
    CALL attach(fid, '/', 'NITS', nits)
    CALL attach(fid, '/', 'LEVELS', levels)
    CALL attach(fid, '/', 'NU1', nu1)
    CALL attach(fid, '/', 'NU2', nu2)
    CALL attach(fid, '/', 'NU0', nu0)
    CALL attach(fid, '/', 'MU', mu)
    CALL creatg(fid, '/mglevels')
    DO l=1,levels
       WRITE(dsname,'("/mglevels/level.",i2.2)') l
       CALL creatg(fid, TRIM(dsname))
       CALL putmat(fid, TRIM(dsname)//'/mata', grids(l)%mata)
       IF(l.GT.1) THEN
          CALL putmat(fid, TRIM(dsname)//'/matpx', grids(l)%matp(1))
          CALL putmat(fid, TRIM(dsname)//'/matpy', grids(l)%matp(2))
       END IF
       CALL putarr(fid, TRIM(dsname)//'/x', grids(l)%x)
       CALL putarr(fid, TRIM(dsname)//'/y', grids(l)%y)
       CALL putarr(fid, TRIM(dsname)//'/f', grids(l)%f)
       CALL putarr(fid, TRIM(dsname)//'/v', grids(l)%v)
       CALL putarr(fid, TRIM(dsname)//'/f1d', grids(l)%f1d)
       CALL putarr(fid, TRIM(dsname)//'/v1d', grids(l)%v1d)
    END DO
!
!  Solutions at finest grid
!
    CALL creatg(fid, '/solutions')
    CALL putarr(fid, '/solutions/xg', grids(1)%x)
    CALL putarr(fid, '/solutions/yg', grids(1)%y)
    CALL putarr(fid, '/solutions/anal', sol_anal_grid)
    CALL putarr(fid, '/solutions/calc', sol_calc_grid)
    CALL putarr(fid, '/solutions/direct', sol_direct_grid)
!
    CALL creatg(fid, '/Iterations')
    CALL putarr(fid, '/Iterations/residues', resid(0:nits))
    CALL putarr(fid, '/Iterations/disc_errors', errdisc(0:nits))
!
    CALL closef(fid)
  END SUBROUTINE h5file
!+++
END PROGRAM