!>
!> @file poisson_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/>.
!>
!> @author
!> (in alphabetical order)
!> @author Trach-Minh Tran <trach-minh.tran@epfl.ch>
!>
MODULE poisson_mod
  IMPLICIT NONE
!
!  Dirichlet BC encapsulated in a derived datatype
!
  TYPE dirich
     INTEGER :: meth, mbess, n1, n2, nidbas1, nidbas2
     INTEGER :: i0, i1
     DOUBLE PRECISION, POINTER :: amat(:,:) => NULL()
     DOUBLE PRECISION, POINTER :: g(:) => NULL()
  END TYPE dirich
!
CONTAINS
  SUBROUTINE disrhs(mbess, spl, rhs)
!
!   Assembly the RHS using 2d spline spl
!
    USE bsplines
    INTEGER, INTENT(in) :: mbess
    TYPE(spline2d), INTENT(in) :: spl
    DOUBLE PRECISION, INTENT(out) :: rhs(:)
!
!   The RHS is 0 
!
    rhs(:) = 0.0d0
  END SUBROUTINE disrhs
!
  SUBROUTINE ibcmat(mat, bc)
!
!   Apply BC on matrix
!
    USE matrix
    TYPE(pbmat), INTENT(inout) :: mat
    TYPE(dirich) :: bc
    INTEGER :: ny
    INTEGER :: kl, ku, nrank, i, j, k
    DOUBLE PRECISION, ALLOCATABLE :: zsum(:), arr(:)
    INTEGER :: i0, i1, ii
!===========================================================================
!              1.0 Prologue
!
    ku = mat%ku
    kl = ku
    nrank = mat%rank
    ny = bc%n2
!===========================================================================
!              2.0 Unicity at the axis
!
!   The vertical sum on the NY-th row
!
    ALLOCATE(zsum(nrank))
    ALLOCATE(arr(nrank))
    zsum = 0.0d0
    DO i=1,ny
       arr = 0.0d0
       CALL getrow(mat, i, arr)
       DO j=1,ny+ku
          zsum(j) = zsum(j) + arr(j)
       END DO
    END DO
!
    zsum(ny) = SUM(zsum(1:ny))   ! using symmetry
    CALL putrow(mat, ny, zsum)
    DEALLOCATE(zsum)
!
!   The away operator
!
    DO i = 1,ny-1
       arr = 0.0d0; arr(i) = 1.0d0
       CALL putrow(mat, i, arr)     
    END DO
    DEALLOCATE(arr)
!===========================================================================
!              3.0 Dirichlet on right boundary
!
!!$    i0 = nrank - ku
!!$    i1 = nrank - ny
    i0 = (bc%n1-1)*bc%n2 + 1
    i1 = nrank - bc%n2
    bc%i0 = i0
    bc%i1 = i1
!
    IF(ASSOCIATED(bc%amat)) DEALLOCATE(bc%amat)
    IF(ASSOCIATED(bc%g)) DEALLOCATE(bc%g)
    ALLOCATE(bc%amat(i0:i1,ny))
    ALLOCATE(bc%g(ny))
!
    WRITE(*,'(/a,2i6)') 'IBCMAT: i0, i1 =', i0, i1
!
!   Extract and save the last ny columns of matrix
!
    ALLOCATE(arr(nrank))
    DO k=1,ny
       j = nrank-ny+k
       CALL getcol(mat, j, arr)
       bc%amat(i0:i1,k) = arr(i0:i1)
       IF( ANY(arr(1:i0-1) .NE. 0.0d0) ) THEN
          WRITE(*,'(a,i4)') 'i0 is underestimated for j =', j
       END IF
    END DO
!
!   The away operator
!
    DO k=1,ny
       j = nrank-ny+k
       arr = 0.0d0; arr(j) = 1.0d0
       CALL putrow(mat, j, arr)     
    END DO
!
    DEALLOCATE(arr)
!
  END SUBROUTINE ibcmat
!+++
  SUBROUTINE ibcrhs(rhs, ygrid, bc)
!
!   Apply BC on RHS
!
    DOUBLE PRECISION, INTENT(inout) :: rhs(:)
    DOUBLE PRECISION, INTENT(in)    :: ygrid(:)
    TYPE(dirich)                    :: bc
!
    INTEGER :: nrank, ny, m, i0, i1
    DOUBLE PRECISION :: zsum
!===========================================================================
!              1.0 Prologue
!
    nrank = SIZE(rhs,1)
    ny = bc%n2
    m = bc%mbess
!===========================================================================
!              2.0 Unicity at the axis
!
!   The vertical sum
!
    zsum = SUM(rhs(1:ny))
    rhs(ny) = zsum
    rhs(1:ny-1) = 0.0d0
!===========================================================================
!              3.0 Dirichlet on right boundary
!
!   Get spline coefs at boundary r=1
!
    SELECT CASE (bc%meth)
    CASE(1)
       CALL dirich_interp(ygrid, bc, frhs)
    CASE(2)
       CALL dirich_minres(ygrid, bc, frhs)
    END SELECT
!
!   Modify RHS
!
    i0 = bc%i0
    i1 = bc%i1
    rhs(i0:i1) = rhs(i0:i1) - MATMUL(bc%amat, bc%g)
    rhs(i1+1:nrank) = bc%g(1:ny)
  CONTAINS
    DOUBLE PRECISION FUNCTION frhs(x)
      DOUBLE PRECISION, INTENT(in) :: x
      frhs = COS(m*x)
    END FUNCTION frhs
  END SUBROUTINE ibcrhs
!++++
  SUBROUTINE coefeq(x, y, idt, idw, c)
    DOUBLE PRECISION, INTENT(in) :: x, y
    INTEGER, INTENT(out) :: idt(:,:), idw(:,:)
    DOUBLE PRECISION, INTENT(out) :: c(:)
!
! Weak form = Int(x*dw/dx*dt/dx + 1/x*dw/dy*dt/dy)dxdy
!
    c(1) = x        ! 
    idt(1,1) = 1
    idt(1,2) = 0
    idw(1,1) = 1
    idw(1,2) = 0
!
    c(2) = 1.d0/x
    idt(2,1) = 0
    idt(2,2) = 1
    idw(2,1) = 0
    idw(2,2) = 1
  END SUBROUTINE coefeq
!++++
  SUBROUTINE dirich_interp(ygrid, bc, frhs)
!
!   Dirichlet BC by interpolation
!
    USE bsplines
    DOUBLE PRECISION, INTENT(in) :: ygrid(:)
    TYPE(dirich)                 :: bc
    INTERFACE
       DOUBLE PRECISION FUNCTION frhs(x)
         DOUBLE PRECISION, INTENT(in) :: x
       END FUNCTION frhs
    END INTERFACE
!
    INTEGER :: nidbas, dim, n2, i
    DOUBLE PRECISION :: shifty
    DOUBLE PRECISION :: gval(SIZE(ygrid))
    DOUBLE PRECISION, ALLOCATABLE :: coefs(:)
    DOUBLE PRECISION :: ygrid_interp(SIZE(ygrid))
    TYPE(spline1d) :: spl_interp
!    
    nidbas = bc%nidbas2
    n2 = bc%n2
!
    IF(MODULO(nidbas,2) .EQ. 0 ) THEN
       shifty = 0.5d0*(ygrid(2)-ygrid(1))
       ygrid_interp(:) = ygrid(:) + shifty
    ELSE
       ygrid_interp(:) = ygrid(:)
    END IF
    CALL set_splcoef(nidbas, ygrid_interp, spl_interp, period=.TRUE.)
    CALL get_dim(spl_interp, dim)
    ALLOCATE(coefs(dim))
!
    DO i=1,SIZE(ygrid)
       gval(i) = frhs(ygrid_interp(i))
    END DO
    CALL get_splcoef(spl_interp, gval, coefs)
!
!   Store spline coefs in bc
!
    bc%g(1:n2) = coefs(1:n2)
!
    DEALLOCATE(coefs)
    CALL destroy_sp(spl_interp)
  END SUBROUTINE dirich_interp
!++++
  SUBROUTINE dirich_minres(xgrid, bc, frhs)
!
!   Dirichlet BC by minimization of residual
!
    USE bsplines
    USE matrix
    USE conmat_mod
    DOUBLE PRECISION, INTENT(in) :: xgrid(:)
    TYPE(dirich)                 :: bc
    INTERFACE
       DOUBLE PRECISION FUNCTION frhs(x)
         DOUBLE PRECISION, INTENT(in) :: x
       END FUNCTION frhs
    END INTERFACE
!
    INTEGER :: nx, nidbas, ngauss, kl, ku
    TYPE(periodic_mat) :: mass_mat
    TYPE(spline1d) :: spl
!
    nidbas = bc%nidbas2
    ngauss = nidbas+1
    nx = bc%n2
    kl = nidbas
    ku = kl
!    
    CALL set_spline(nidbas, ngauss, xgrid, spl, period=.TRUE.)   
    CALL init(kl, ku, nx, 1, mass_mat)
    CALL conmat(spl, mass_mat, coefeq_mass)
    CALL conrhs(spl, bc%g, frhs)
    CALL factor(mass_mat)
    CALL bsolve(mass_mat, bc%g)
!
    CALL destroy(mass_mat)
    CALL destroy_sp(spl)
!
  CONTAINS
    SUBROUTINE coefeq_mass(x, idt, idw, c)
      DOUBLE PRECISION, INTENT(in) :: x
      INTEGER, INTENT(out) :: idt(:), idw(:)
      DOUBLE PRECISION, INTENT(out) :: c(:)
      c(1) = 1.0d0
      idt(1) = 0
      idw(1) = 0
    END SUBROUTINE coefeq_mass
  END SUBROUTINE dirich_minres
!++++
  DOUBLE PRECISION FUNCTION err2_norm(spl, jder, mbess, fexact)
!
!  Compute error L2 norm unsing Gauss points
!
    USE bsplines
    TYPE(spline2d) :: spl
    INTEGER, INTENT(in) :: jder(:)
    INTEGER, INTENT(in) :: mbess
    INTERFACE
       DOUBLE PRECISION FUNCTION fexact(m,x,y)
         INTEGER, INTENT(in)          :: m
         DOUBLE PRECISION, INTENT(in) :: x, y
       END FUNCTION fexact
    END INTERFACE
!
    DOUBLE PRECISION, POINTER :: xg1(:,:), xg2(:,:), wg1(:,:), wg2(:,:)
    INTEGER :: i1, ig1, n1, nidbas1, ndim1, ng1
    INTEGER :: i2, ig2, n2, nidbas2, ndim2, ng2
    DOUBLE PRECISION :: contrib
    DOUBLE PRECISION, ALLOCATABLE :: x(:), y(:), sol(:,:)
!
!   Gauss points and weights on all intervals
!
    CALL get_dim(spl%sp1, ndim1, n1, nidbas1)
    CALL get_dim(spl%sp2, ndim2, n2, nidbas2)
    xg1 => spl%sp1%gausx  ! xg1(ng1,n1)
    wg1 => spl%sp1%gausw  ! wg1(ng1,n1)
    ng1 = SIZE(xg1,1)
    xg2 => spl%sp2%gausx
    wg2 => spl%sp2%gausw
    ng2 = SIZE(xg2,1)
!
    err2_norm = 0.0d0
    ALLOCATE(x(ng1), y(ng2))
    ALLOCATE(sol(ng1,ng2))
    DO i1=1,n1                
       x=xg1(:,i1)
       DO i2=1,n2
          y=xg2(:,i2)
          CALL gridval(spl, x, y, sol, jder)
          DO ig1=1,ng1            
             DO ig2=1,ng2      
                contrib = wg1(ig1,i1)*wg2(ig2,i2)*(sol(ig1,ig2) - &
                     &    fexact(mbess,x(ig1),y(ig2)))**2
                err2_norm = err2_norm +  x(ig1)*contrib !use same inner-product in weak-form
             END DO
          END DO
       END DO
    END DO
    DEALLOCATE(x)
    DEALLOCATE(y)
    DEALLOCATE(sol)
    err2_norm = SQRT(err2_norm)
  END FUNCTION err2_norm
END MODULE poisson_mod