!>
!> @file poisson_petsc.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
  USE petsc_bsplines
  IMPLICIT NONE
  TYPE(petsc_mat) :: amat
  DOUBLE PRECISION, ALLOCATABLE :: rhs(:), sol(:), arow(:)
  INTEGER :: nx=5, ny=4, ntrials=10
  INTEGER :: nitmax=10000, nits
  DOUBLE PRECISION :: rtol=1.e-9
  INTEGER :: n, nnz, nnz_loc
  INTEGER :: i, j, irow, jcol
  INTEGER :: ierr, me, npes, istart, iend
  DOUBLE PRECISION :: t0
  INTEGER :: ncols, cols(5)   !  Max nnz by row   .LE. 5
!
  CHARACTER(len=128) :: matfile='mat.dat', rhsfile='rhs.dat'
  LOGICAL :: file_exist
!
  NAMELIST /newrun/ nx, ny, nitmax, rtol, matfile, rhsfile
!
  CALL mpi_init(ierr)
  CALL mpi_comm_rank(MPI_COMM_WORLD, me, ierr)
  CALL mpi_comm_size(MPI_COMM_WORLD, npes, ierr)
!
  IF(me.EQ.0) THEN
     READ(*, newrun)
     WRITE(*, newrun)
     WRITE(*,'(a,i6)') 'npes =', npes
  END IF
  CALL mpi_bcast(nx, 1, MPI_INTEGER, 0, MPI_COMM_WORLD, ierr)
  CALL mpi_bcast(ny, 1, MPI_INTEGER, 0, MPI_COMM_WORLD, ierr)
  CALL mpi_bcast(nitmax, 1, MPI_INTEGER, 0, MPI_COMM_WORLD, ierr)
  CALL mpi_bcast(rtol, 1, MPI_DOUBLE_PRECISION, 0, MPI_COMM_WORLD, ierr)
  CALL mpi_bcast(matfile, 128, MPI_CHARACTER, 0, MPI_COMM_WORLD, ierr)
  CALL mpi_bcast(rhsfile, 128, MPI_CHARACTER, 0, MPI_COMM_WORLD, ierr)
!
  n = nx*ny  ! Rank of the matrix
!
!   Initialize matrix
!
  CALL init(n, 1, amat, comm=MPI_COMM_WORLD)
  istart = amat%istart
  iend   = amat%iend
!!$  WRITE(*,'(a,i3.3,a,3i6)') 'PE', me, ': istart, iend', istart, iend
!
!
  INQUIRE(file=TRIM(matfile), exist=file_exist)
!
  IF( file_exist ) THEN
     CALL mpi_barrier(MPI_COMM_WORLD,ierr)
     t0 = mpi_wtime()
     CALL load_mat(amat, matfile)
     CALL mpi_barrier(MPI_COMM_WORLD,ierr)
     IF(me.EQ.0) WRITE(*,'(a,1pe12.3)') 'Mat. loading time (s)', mpi_wtime()-t0
  ELSE
!
!  Construct the matrix
!
     CALL mpi_barrier(MPI_COMM_WORLD,ierr)
     t0 = mpi_wtime()
     ALLOCATE(arow(5))   !  Max nnz per row .LE. 5
     DO j=1,ny
        DO i=1,nx
           irow = numb(i,j)
           IF( irow.GE.istart .AND. irow.LE.iend) THEN
              ncols = 1; cols(ncols) = irow; arow(ncols) = 4.0d0
              IF(i.GT.1)  THEN 
                 ncols = ncols+1
                 cols(ncols) = numb(i-1,j); arow(ncols) = -1.0d0
              END IF
              IF(i.LT.nx) THEN
                 ncols = ncols+1
                 cols(ncols) = numb(i+1,j); arow(ncols) = -1.0d0
              END IF
              IF(j.GT.1)  THEN
                 ncols = ncols+1
                 cols(ncols) = numb(i,j-1); arow(ncols) = -1.0d0
              END IF
              IF(j.LT.ny) THEN 
                 ncols = ncols+1
                 cols(ncols) = numb(i,j+1); arow(ncols) = -1.0d0
              END IF
              CALL putrow(amat, irow, arow(1:ncols), cols(1:ncols))
           END IF
        END DO
     END DO
     DEALLOCATE(arow)
     CALL mpi_barrier(MPI_COMM_WORLD,ierr)
     IF(me.EQ.0) WRITE(*,'(a,1pe12.3)') 'Mat. construction time (s)', mpi_wtime()-t0
!
!   Convert to PETSC mat
!
     CALL mpi_barrier(MPI_COMM_WORLD,ierr)
     t0=mpi_wtime()
     CALL to_mat(amat)
     CALL mpi_barrier(MPI_COMM_WORLD,ierr)
     IF(me.EQ.0) WRITE(*,'(a,1pe12.3)') 'Mat. conversion time (s)', mpi_wtime()-t0
!
     CALL save_mat(amat, matfile)
  END IF
!
!  Matrix size and partition could have changed after loading from file!
!
  n = amat%rank
  istart = amat%istart
  iend   = amat%iend
!
  nnz_loc = get_count(amat)
  CALL mpi_reduce(nnz_loc, nnz, 1, MPI_INTEGER, mpi_sum, 0, MPI_COMM_WORLD, ierr)
  IF(npes.LE.4) THEN
     WRITE(*,'(a,i3.3,a,3i6)') 'PE', me, ': istart, iend (after), nloc, nnz_loc', &
          &   istart, iend, iend-istart+1, nnz_loc
  END IF
!
!  Construct or read  RHS
!
  CALL mpi_barrier(MPI_COMM_WORLD,ierr)
  t0 = mpi_wtime()
  ALLOCATE(rhs(n))
  INQUIRE(file=TRIM(rhsfile), exist=file_exist)
  IF( file_exist ) THEN
     OPEN(unit=99, file=TRIM(rhsfile), status='old', form='unformatted')
     READ(99) rhs(1:n)
     CLOSE(99)
     CALL mpi_barrier(MPI_COMM_WORLD,ierr)
     IF(me.EQ.0) THEN
        WRITE(*,'(a,1pe12.3)') 'RHS read time (s)', mpi_wtime()-t0
     END IF
  ELSE
     rhs = 0.0d0
     ALLOCATE(arow(n))
     DO i=istart, iend 
        arow = 0.0d0
        CALL getrow(amat, i, arow)
        rhs(i) = SUM(arow)   ! => the exact solution is 1
     END DO
     arow = rhs
     CALL mpi_allreduce(arow, rhs, n, MPI_DOUBLE_PRECISION, MPI_SUM, MPI_COMM_WORLD, ierr)
     DEALLOCATE(arow)
     IF( me.EQ.0 ) THEN   ! All processes have the gobla rhs
        OPEN(unit=99, file=TRIM(rhsfile), status='new', form='unformatted')
        WRITE(99) rhs(1:n)
        CLOSE(99)
     END IF
     CALL mpi_barrier(MPI_COMM_WORLD,ierr)
     IF(me.EQ.0) THEN
        WRITE(*,'(a,1pe12.3)') 'RHS construction time (s)', mpi_wtime()-t0
     END IF
  END IF
  CLOSE(99)
  CALL mpi_barrier(MPI_COMM_WORLD,ierr)
!
!  Back solve
! 
  ALLOCATE(sol(n))
!
  CALL mpi_barrier(MPI_COMM_WORLD,ierr)
  sol = 0.0d0
  t0=mpi_wtime()
  CALL bsolve(amat, rhs, sol, rtol, nitmax, nits)
  CALL mpi_barrier(MPI_COMM_WORLD,ierr)
  IF(me.EQ.0) THEN
     WRITE(*,'(a,1pe12.3,i8,1pe12.3)') 'Error, nits, solve time(s)', &
          &   MAXVAL(ABS(sol-1.0d0)), nits, mpi_wtime()-t0
  END IF
  CALL mpi_barrier(MPI_COMM_WORLD,ierr)
  t0=mpi_wtime()
  DO i=1,ntrials
     sol = 0.0d0
     CALL bsolve(amat, rhs, sol, rtol, nitmax, nits)
  END DO
  CALL mpi_barrier(MPI_COMM_WORLD,ierr)
  IF(me.EQ.0) THEN
     WRITE(*,'(a,1pe12.3,i8,1pe12.3)') 'Error, nits, solve time(s)', &
          &   MAXVAL(ABS(sol-1.0d0)), nits, (mpi_wtime()-t0)/REAL(ntrials)
  END IF
!
!   Clean up
!
  DEALLOCATE(rhs)
  DEALLOCATE(sol)
  CALL destroy(amat)
  CALL PetscFinalize(ierr)
  CALL mpi_finalize(ierr)
CONTAINS
  INTEGER FUNCTION numb(i,j)
!
!  One-dimensional numbering
!  Number first x then y
!
    INTEGER, INTENT(in) :: i, j
    numb = (j-1)*nx + i
  END FUNCTION numb
END PROGRAM main