!>
!> @file gvector_mod.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>
!>
MODULE gvector
!
!    Implementation of 2D vectors with arbitrary
!    vector bounds and ghost cell width.
!
!  T.M. Tran, CRPP-EPFL
!  September 2013
!
  USE iso_fortran_env, ONLY : rkind => real64
  IMPLICIT NONE
  PRIVATE
  PUBLIC :: gvector_2d, disp, norm2, &
       &    OPERATOR(+), OPERATOR(-), OPERATOR(*), &
       &    ASSIGNMENT(=)

  TYPE gvector_2d
     INTEGER, DIMENSION(2)    :: s, e ! vector internal bounds
     INTEGER, DIMENSION(2)    :: g    ! ghost cell widths
     REAL(rkind), ALLOCATABLE :: val(:,:)
  END TYPE gvector_2d

  INTERFACE gvector_2d
     MODULE PROCEDURE constructor
  END INTERFACE gvector_2d
  INTERFACE OPERATOR(+)
     MODULE PROCEDURE add_scal
     MODULE PROCEDURE add_vec
  END INTERFACE OPERATOR(+)
  INTERFACE OPERATOR(-)
     MODULE PROCEDURE minus_vec
     MODULE PROCEDURE substract_vec
  END INTERFACE OPERATOR(-)
  INTERFACE OPERATOR(*)
     MODULE PROCEDURE scale_left
     MODULE PROCEDURE scale_right
  END INTERFACE OPERATOR(*)
  INTERFACE ASSIGNMENT(=)
     MODULE PROCEDURE from_scal
     MODULE PROCEDURE from_vec
  END INTERFACE ASSIGNMENT(=)
  INTERFACE norm2
     MODULE PROCEDURE norm2_gvector_2d
     MODULE PROCEDURE norm2_root_g_2d
     MODULE PROCEDURE norm2_all_g_2d
  END INTERFACE norm2

CONTAINS
!=======================================================================
  FUNCTION constructor(s, e, g) RESULT(res)
    INTEGER, INTENT(in)           :: s(2), e(2)
    INTEGER, OPTIONAL, INTENT(in) :: g(2)
    TYPE(gvector_2d)              :: res
    INTEGER :: lb(2), ub(2)
    res%g= 0
    IF(PRESENT(g)) res%g = g
    res%s = s
    res%e = e
    lb = res%s - res%g
    ub = res%e + res%g
    ALLOCATE(res%val(lb(1):ub(1),lb(2):ub(2)))
!
! Initialize to 0 on all ghost cells
    res%val(lb(1):s(1)-1,:) = 0._rkind
    res%val(e(1)+1:ub(1),:) = 0._rkind
    res%val(:,lb(2):s(2)-1) = 0._rkind
    res%val(:,e(2)+1:ub(2)) = 0._rkind
  END FUNCTION constructor
!=======================================================================
  FUNCTION add_vec(lhs, rhs) RESULT(res)
    TYPE(gvector_2d), INTENT(in) :: lhs, rhs
    TYPE(gvector_2d)             :: res
    res = gvector_2d(lhs%s, lhs%e, lhs%g)
    res%val(res%s(1):res%e(1),res%s(2):res%e(2)) = &
         &  lhs%val(res%s(1):res%e(1),res%s(2):res%e(2)) + &
         &  rhs%val(res%s(1):res%e(1),res%s(2):res%e(2))
  END FUNCTION add_vec
!=======================================================================
  FUNCTION add_scal(lhs, rhs) RESULT(res)
    TYPE(gvector_2d), INTENT(in) :: lhs
    REAL(rkind), INTENT(in)      :: rhs
    TYPE(gvector_2d)             :: res
    res = gvector_2d(lhs%s, lhs%e, lhs%g)
    res%val(res%s(1):res%e(1),res%s(2):res%e(2)) = &
         &  lhs%val(res%s(1):res%e(1),res%s(2):res%e(2)) + rhs
  END FUNCTION add_scal
!=======================================================================
  FUNCTION minus_vec(this) RESULT(res)
    TYPE(gvector_2d), INTENT(in) :: this
    TYPE(gvector_2d)             :: res
    res = gvector_2d(this%s, this%e, this%g)
    res%val(res%s(1):res%e(1),res%s(2):res%e(2)) = &
         &  -this%val(res%s(1):res%e(1),res%s(2):res%e(2))
  END FUNCTION minus_vec
!=======================================================================
  FUNCTION substract_vec(lhs, rhs) RESULT(res)
    TYPE(gvector_2d), INTENT(in) :: lhs, rhs
    TYPE(gvector_2d)             :: res
    res = gvector_2d(lhs%s, lhs%e, lhs%g)
    res%val(res%s(1):res%e(1),res%s(2):res%e(2)) = &
         &  lhs%val(res%s(1):res%e(1),res%s(2):res%e(2)) - &
         &  rhs%val(res%s(1):res%e(1),res%s(2):res%e(2))
  END FUNCTION substract_vec
!=======================================================================
  FUNCTION scale_left(lhs, rhs) RESULT(res)
    REAL(rkind), INTENT(in)      :: lhs
    TYPE(gvector_2d), INTENT(in) :: rhs
    TYPE(gvector_2d)             :: res
    res = gvector_2d(rhs%s, rhs%e, rhs%g)
    res%val(res%s(1):res%e(1),res%s(2):res%e(2)) = &
         &  lhs * rhs%val(res%s(1):res%e(1),res%s(2):res%e(2))
  END FUNCTION scale_left
!=======================================================================
  FUNCTION scale_right(lhs, rhs) RESULT(res)
    TYPE(gvector_2d), INTENT(in) :: lhs
    REAL(rkind), INTENT(in)      :: rhs
    TYPE(gvector_2d)             :: res
    res = gvector_2d(lhs%s, lhs%e, lhs%g)
    res%val(res%s(1):res%e(1),res%s(2):res%e(2)) = &
         &  lhs%val(res%s(1):res%e(1),res%s(2):res%e(2)) * rhs
  END FUNCTION scale_right
!=======================================================================

  SUBROUTINE from_vec(lhs, rhs)
    TYPE(gvector_2d), INTENT(inout) :: lhs
    REAL(rkind), INTENT(in)         :: rhs(:,:)
    INTEGER :: n(2)
    n = lhs%e - lhs%s + 1
    IF(SIZE(rhs,1).NE.n(1) .OR. SIZE(rhs,2).NE.n(2)) THEN
       PRINT*, 'from_vec: sizes of rhs and lhs not equal!'
       STOP
    END IF
    lhs%val(lhs%s(1):lhs%e(1),lhs%s(2):lhs%e(2)) = rhs(:,:)
  END SUBROUTINE from_vec
!=======================================================================

  SUBROUTINE from_scal(lhs, rhs)
    TYPE(gvector_2d), INTENT(inout) :: lhs
    REAL(rkind), INTENT(in)         :: rhs
    lhs%val(lhs%s(1):lhs%e(1),lhs%s(2):lhs%e(2)) = rhs
  END SUBROUTINE from_scal
!=======================================================================
  SUBROUTINE disp(str,this)
    CHARACTER(len=*), INTENT(in) :: str
    TYPE(gvector_2d), INTENT(in) :: this
    INTEGER :: i
    WRITE(*,'(/a,3(" (",i0,",",i0,") "))') str//': s, e, g =',&
         &   this%s, this%e, this%g
    DO i=LBOUND(this%val,1),UBOUND(this%val,1)
       WRITE(*,'(10(1pe11.3))') (this%val(i,:))
    END DO
  END SUBROUTINE disp
!=======================================================================

  FUNCTION norm2_gvector_2d(this) RESULT(res)
    TYPE(gvector_2d), INTENT(in) :: this
    REAL(rkind)                  :: res
    res = NORM2( this%val(this%s(1):this%e(1), &
         &  this%s(2):this%e(2)) )
  END FUNCTION norm2_gvector_2d
!=======================================================================
  FUNCTION norm2_root_g_2d(x, comm, root) RESULT(res)
!
!  Vector norm of 2d distributed array with ghost cells
!
    USE mpi
    TYPE(gvector_2d), INTENT(in) :: x
    INTEGER, INTENT(in)          :: comm
    INTEGER, INTENT(in)          :: root
    REAL(rkind)                  :: res
    INTEGER, PARAMETER       :: ndim=2
    INTEGER, DIMENSION(ndim) :: s, e
    REAL(rkind)              :: res_loc
    INTEGER                  :: me, ierr
!
    CALL mpi_comm_rank(comm, me, ierr)
    s = x%s
    e = x%e
    res_loc = SUM(x%val(s(1):e(1),s(2):e(2))**2)
    res = 0.0
    CALL mpi_reduce(res_loc, res, 1, MPI_DOUBLE_PRECISION, MPI_SUM,&
         &          root, comm, ierr)
    IF(me.EQ.root) res = SQRT(res)    
  END FUNCTION norm2_root_g_2d
!=======================================================================
  FUNCTION norm2_all_g_2d(x, comm) RESULT(res)
!
!  Vector norm of 2d distributed array with ghost cells
!
    USE mpi
    TYPE(gvector_2d), INTENT(in) :: x
    INTEGER, INTENT(in)          :: comm
    REAL(rkind)                  :: res
    INTEGER, PARAMETER       :: ndim=2
    INTEGER, DIMENSION(ndim) :: s, e
    REAL(rkind)              :: res_loc
    INTEGER                  :: ierr
!
    s = x%s
    e = x%e
    res_loc = SUM(x%val(s(1):e(1),s(2):e(2))**2)
    CALL mpi_allreduce(res_loc, res, 1, MPI_DOUBLE_PRECISION, MPI_SUM, &
         &          comm, ierr)
    res = SQRT(res)    
  END FUNCTION norm2_all_g_2d
!=======================================================================
END MODULE gvector