!>
!> @file tcdsmat_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 tcdsmat_mod
  IMPLICIT NONE
!
  INTERFACE
     SUBROUTINE meshdist(coefs, x, nx)
       DOUBLE PRECISION, INTENT(in) :: coefs(5)
       INTEGER, INTENT(iN) :: nx
       DOUBLE PRECISION, INTENT(inout) :: x(0:nx)
     END SUBROUTINE meshdist
  END INTERFACE
!
  INTERFACE norm2
     MODULE PROCEDURE norm2_1d, norm2_2d
  END INTERFACE
CONTAINS
!++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
  SUBROUTINE mstruct(p, n, rank, dist)
!
!  It is assumed that:
!    . 2nd dimension is number first
!    . 2nd dimension is periodic
!
    INTEGER, INTENT(in)  :: p(2), n(2)
    INTEGER, INTENT(out) :: rank
    INTEGER, ALLOCATABLE :: dist(:)
!
    INTEGER, ALLOCATABLE :: pvect(:,:)
    INTEGER              :: kl, ku, i
!
    rank = (n(1)+p(1))*n(2)    ! Rank of the FE matrix
    ku = (p(1)+1)*(2*p(2)+1)-1
    kl = ku
    ALLOCATE(pvect(0:2*p(2),0:p(1)))
    IF( ALLOCATED(dist)) DEALLOCATE(dist)
    ALLOCATE(dist(-kl:ku))
!
!   Upper (North) points
    pvect(0:p(2),0) = (/(i,i=0,p(2))/)
!
!   Lower (South) points and periodicity of 2nd dim.
    DO i=1,p(2)
       pvect(p(2)+i,0) = n(2)-pvect(p(2)-i+1,0)
    END DO
!
!  Shift by N2 for points on the right (West) side
    DO i=1,p(1)
       pvect(:,i) = pvect(:,i-1)+n(2)
    END DO
!
!   Super-diagonals including the diagonal
    dist(0:ku) = RESHAPE(pvect, (/ku+1/))
!
!   Sub-diagonals
    DO i=-1,-kl,-1
       dist(i) = -dist(-i)
    END DO
!
    DEALLOCATE(pvect)
  END SUBROUTINE mstruct
!++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
  SUBROUTINE dismat(spl, mat)
!
!   Assembly of FE matrix mat using spline spl
!
    USE bsplines
    USE cds
!
    TYPE(spline2d), INTENT(in)   :: spl
    TYPE(cds_mat), INTENT(inout) :: mat
!
    INTEGER :: n1, nidbas1, ndim1, ng1
    INTEGER :: n2, nidbas2, ndim2, ng2
    INTEGER :: i, j, ig1, ig2
    INTEGER :: iterm, iw1, iw2, igw1, igw2, it1, it2, igt1, igt2, irow, jcol
    DOUBLE PRECISION, ALLOCATABLE :: xg1(:), wg1(:), fun1(:,:)
    DOUBLE PRECISION, ALLOCATABLE :: xg2(:), wg2(:), fun2(:,:)
    DOUBLE PRECISION:: contrib
!
    INTEGER :: kterms         ! Number of terms in weak form
    INTEGER, DIMENSION(:,:), ALLOCATABLE :: idert, iderw  ! Derivative order
    DOUBLE PRECISION, DIMENSION(:), ALLOCATABLE  :: coefs ! Terms in weak form
!===========================================================================
!              1.0 Prologue
!
!   Properties of spline space
!
    CALL get_dim(spl%sp1, ndim1, n1, nidbas1)
    CALL get_dim(spl%sp2, ndim2, n2, nidbas2)
    WRITE(*,'(/a, 5i6)') 'n1, nidbas1, ndim1 =', n1, nidbas1, ndim1
    WRITE(*,'(a, 5i6)') 'n2, nidbas2, ndim2 =', n2, nidbas2, ndim2
!
!   Weak form
!
    kterms = mat%nterms
    ALLOCATE(idert(kterms,2), iderw(kterms,2), coefs(kterms))
!
    ALLOCATE(fun1(0:nidbas1,0:1)) ! Spline and 1st derivative
    ALLOCATE(fun2(0:nidbas2,0:1)) ! 
!
!   Gauss quadature
!
    CALL get_gauss(spl%sp1, ng1)
    CALL get_gauss(spl%sp2, ng2)
    ALLOCATE(xg1(ng1), wg1(ng1))
    ALLOCATE(xg2(ng1), wg2(ng1))
 !===========================================================================
!              2.0 Assembly loop
!
    DO i=1,n1
       CALL get_gauss(spl%sp1, ng1, i, xg1, wg1)
       DO ig1=1,ng1
          CALL basfun(xg1(ig1), spl%sp1, fun1, i)
          DO j=1,n2
             CALL get_gauss(spl%sp2, ng2, j, xg2, wg2)
             DO ig2=1,ng2
                CALL basfun(xg2(ig2), spl%sp2, fun2, j)
                CALL coefeq(xg1(ig1), xg2(ig2), idert, iderw, coefs)
                DO iterm=1,kterms
                   DO iw1=0,nidbas1  ! Weight function in dir 1
                      igw1 = i+iw1
                      DO iw2=0,nidbas2  ! Weight function in dir 2
                         igw2 = MODULO(j+iw2-1, n2) + 1
                         irow = igw2 + (igw1-1)*n2
                         DO it1=0,nidbas1  ! Test function in dir 1
                            igt1 = i+it1
                            DO it2=0,nidbas2  ! Test function in dir 2
                               igt2 = MODULO(j+it2-1, n2) + 1
                               jcol = igt2 + (igt1-1)*n2
                               contrib = fun1(iw1,iderw(iterm,1)) * &
                                    &    fun2(iw2,iderw(iterm,2)) * &
                                    &    coefs(iterm) *             &
                                    &    fun2(it2,idert(iterm,2)) * &
                                    &    fun1(it1,idert(iterm,1)) * &
                                    &    wg1(ig1) * wg2(ig2)
                               CALL updtmat(mat, irow, jcol, contrib)
                            END DO
                         END DO
                      END DO
                   END DO
                END DO
             END DO
          END DO
       END DO
    END DO
!===========================================================================
!              9.0  Epilogue
!
    DEALLOCATE(xg1, wg1, fun1)
    DEALLOCATE(xg2, wg2, fun2)
    DEALLOCATE(idert, iderw, coefs)
!
  CONTAINS
    SUBROUTINE coefeq(x, y, idt, idw, c)
      DOUBLE PRECISION, INTENT(in) :: x, y
      INTEGER, INTENT(out) :: idt(:,:), idw(:,:)
      DOUBLE PRECISION, INTENT(out) :: c(SIZE(idt,1))
!
! 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
  END SUBROUTINE dismat
!++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
  SUBROUTINE ibcmat(mat, ny, px)
!
!   Apply BC on matrix
!
    USE cds
    IMPLICIT NONE
    TYPE(cds_mat), INTENT(inout) :: mat
    INTEGER, INTENT(in)          :: ny, px
!
    INTEGER :: kl, ku, n, bw0, i, j
    DOUBLE PRECISION :: zsum(mat%rank), arr(mat%rank)
!===========================================================================
!              1.0 Prologue
!
    kl = mat%kl
    ku = mat%ku
    n = mat%rank
    mat%ny = ny
!
!   Size of  row ny and column ny
!
    bw0 = SIZE(mat%rowv)
!===========================================================================
!              2.0 Unicity at the axis
!
!   The vertical sum on the NY-th row
!   Store the sums at row ny in MAT%ROWV
!
    zsum(1:n) = 0.0d0
    DO i=1,ny
       CALL getrow(mat, i, arr)
       zsum(1:bw0) = zsum(1:bw0) + arr(1:bw0)
       IF( i .LE. ny ) THEN    ! Clear rows 1:(ny-1)
          arr = 0.0d0; arr(i) = 1.0d0
          CALL putrow(mat, i, arr)
       END IF
    END DO
    mat%rowv(1:bw0) = zsum(1:bw0)   ! row ny
!!$    WRITE(*,'(/a,/(10(f8.3)))') 'rowv', mat%rowv(1:130)
!
!   The horizontal sum on the NY-th column
!   The NY-th row of matrix was stored in mat%rowv
!   Store the sums ar column ny at MAT%COLH
!
    zsum(1:n) = 0.0d0
    DO j=1,ny
       CALL getcol(mat, j, arr)
       zsum(ny) = zsum(ny) + mat%rowv(j)
       zsum(ny+1:bw0) = zsum(ny+1:bw0) + arr(ny+1:bw0)
       IF( j .NE. ny ) THEN    ! Clear columns 1:(ny-1)
          mat%rowv(j) = 0.0d0
          arr = 0.0d0; arr(j) = 1.0d0
          CALL putcol(mat, j, arr)
       END IF
    END DO
    mat%rowv(ny) = 0.0d0   ! Its value is now in mat%colh(ny)
    arr = 0.0d0
    CALL putcol(mat, ny, arr)
    CALL putrow(mat, ny, arr)
    mat%colh(1:bw0) = zsum(1:bw0) ! column ny
!
!   Move the diagonal term from mat%colh back to main diagonal
!
    CALL putele(mat, ny, ny, mat%colh(ny))
    mat%colh(ny) = 0.0d0
!!$    WRITE(*,'(/a,/(10(f8.3)))') 'rowv', mat%rowv(1:130)
!!$    WRITE(*,'(/a,/(10(f8.3)))') 'colh', mat%colh(1:130)
!!$    WRITE(*,'(/a,/(10(f8.3)))') 'colh', mat%val(1:ny,0)
!===========================================================================
!              3.0 Dirichlet on right boundary
!
    DO j = n, n-ny+1, -1
       arr = 0.0d0; arr(j) = 1.0d0
       CALL putcol(mat, j, arr)     
    END DO
!
    DO i = n, n-ny+1, -1
       arr = 0.0d0; arr(i) = 1.0d0
       CALL putrow(mat, i, arr)     
    END DO
!!$    WRITE(*,'(/a,/(10(f8.3)))') 'diag 1', mat%val(n-ny-1:n,1)
!!$    WRITE(*,'(/a,/(10(f8.3)))') 'diag 0', mat%val(n-ny:n,0)
!!$    WRITE(*,'(/a,/(10(f8.3)))') 'diag -1', mat%val(n-ny-1:n,-1)
!===========================================================================
!              9.0  Epilogue
!
  END SUBROUTINE ibcmat
!===========================================================================
  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(:)
!
    INTEGER                       :: n1, nidbas1, ndim1, ng1
    INTEGER                       :: n2, nidbas2, ndim2, ng2
    INTEGER                       :: i, j, ig1, ig2, k1, k2, i1, j2, ij, nrank
    DOUBLE PRECISION, ALLOCATABLE :: xg1(:), wg1(:), fun1(:,:)
    DOUBLE PRECISION, ALLOCATABLE :: xg2(:), wg2(:), fun2(:,:)
    DOUBLE PRECISION:: contrib
!===========================================================================
!              1.0 Prologue
!
!   Properties of spline space
!
    CALL get_dim(spl%sp1, ndim1, n1, nidbas1)
    CALL get_dim(spl%sp2, ndim2, n2, nidbas2) 
!
    ALLOCATE(fun1(0:nidbas1,1)) ! needs only basis functions (no derivatives)
    ALLOCATE(fun2(0:nidbas2,1)) ! needs only basis functions (no derivatives)
!
!   Gauss quadature
!
    CALL get_gauss(spl%sp1, ng1)
    CALL get_gauss(spl%sp2, ng2)
    WRITE(*,'(/a, 2i3)') 'Gauss points and weights, ngauss =', ng1, ng2
    ALLOCATE(xg1(ng1), wg1(ng1))
    ALLOCATE(xg2(ng1), wg2(ng1))
!===========================================================================
!              2.0 Assembly loop
!
    nrank = SIZE(rhs)
    rhs(1:nrank) = 0.0d0
!
    DO i=1,n1
       CALL get_gauss(spl%sp1, ng1, i, xg1, wg1)
       DO ig1=1,ng1
          CALL basfun(xg1(ig1), spl%sp1, fun1, i)
          DO j=1,n2
             CALL get_gauss(spl%sp2, ng2, j, xg2, wg2)
             DO ig2=1,ng2
                CALL basfun(xg2(ig2), spl%sp2, fun2, j)
                contrib = wg1(ig1)*wg2(ig2) * &
                     &    rhseq(xg1(ig1),xg2(ig2), mbess)
                DO k1=0,nidbas1
                   i1 = i+k1
                   DO k2=0,nidbas2
                      j2 = MODULO(j+k2-1,n2) + 1
                      ij = j2 + (i1-1)*n2
                      rhs(ij) = rhs(ij) + contrib*fun1(k1,1)*fun2(k2,1)
                   END DO
                END DO
             END DO
          END DO
       END DO
    END DO
!===========================================================================
!              9.0  Epilogue
!
    DEALLOCATE(xg1, wg1, fun1)
    DEALLOCATE(xg2, wg2, fun2)
!
  CONTAINS
    DOUBLE PRECISION FUNCTION rhseq(x1, x2, m)
      DOUBLE PRECISION, INTENT(in) :: x1, x2
      INTEGER, INTENT(in) :: m
      rhseq = REAL(4*(m+1),8)*x1**(m+1)*COS(REAL(m,8)*x2)
    END FUNCTION rhseq
  END SUBROUTINE disrhs
!++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
  SUBROUTINE ibcrhs(rhs, ny)
!
!   Apply BC on RHS
!
    DOUBLE PRECISION, INTENT(inout) :: rhs(:)
    INTEGER, INTENT(in)              :: ny
!
    INTEGER          :: nrank
    DOUBLE PRECISION :: zsum
!===========================================================================
!              1.0 Prologue
!
    nrank = SIZE(rhs,1)
!===========================================================================
!              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
!
    rhs(nrank-ny+1:nrank) = 0.0d0
  END SUBROUTINE ibcrhs
!===========================================================================
  SUBROUTINE cg(mat, rhs, omega, nitmx, rtolmx, sol, resid, nit, verbose, &
       &        nssor)
!
!   Preconditionned Conjugate Gradient solver
!
    USE cds
    TYPE(cds_mat)                           :: mat
    DOUBLE PRECISION, INTENT(in)            :: rhs(:), omega, rtolmx
    INTEGER, INTENT(in)                     :: nitmx
    DOUBLE PRECISION, INTENT(inout)         :: sol(:)
    DOUBLE PRECISION, OPTIONAL, INTENT(out) :: resid
    INTEGER, OPTIONAL, INTENT(out)          :: nit
    INTEGER, OPTIONAL, INTENT(in)           :: nssor
    LOGICAL, OPTIONAL, INTENT(in)           :: verbose
!
    DOUBLE PRECISION, DIMENSION(SIZE(rhs,1)) :: wr, wz, wp, wq
    DOUBLE PRECISION :: bnrm2, residue, rho0, rho1, alpha, beta
    INTEGER          :: it
!
    bnrm2 = SQRT(DOT_PRODUCT(rhs,rhs))  ! Euclidian norm of RHS
    it = 0
    wr = rhs-vmx(mat,sol)
!
!... Iteration loop  (see fig. 2.5, p.15 of "Templates...")
    DO
       it = it+1
       IF( PRESENT(nssor) ) THEN
          CALL psolve(mat, wz, wr, omega, nssor)
       ELSE
          wz = wr
       END IF
       rho1 = DOT_PRODUCT(wr,wz)
       IF( it .EQ. 1 ) THEN
          wp = wz
       ELSE
          beta = rho1/rho0
          wp = wz + beta*wp
       END IF
       wq = vmx(mat,wp)
       alpha = rho1 / DOT_PRODUCT(wp,wq)
       sol = sol + alpha*wp
       wr = wr - alpha*wq
       residue = SQRT(DOT_PRODUCT(wr,wr)) / bnrm2
       IF( PRESENT(verbose) ) THEN
          IF(verbose) WRITE(*,'(a,i8,1pe12.3)') 'it, resid', it, residue
       END IF
       IF( residue .LE. rtolmx .OR. it .GE. nitmx) EXIT
       rho0 = rho1
    END DO
    IF(PRESENT(resid)) resid = residue
    IF(PRESENT(nit)) nit = it
  END SUBROUTINE cg
!===========================================================================
  SUBROUTINE psolve(mat, x, b, omega, niter_in)
!
!   Preconditionners
!
    USE cds
    TYPE(cds_mat)                 :: mat
    DOUBLE PRECISION, INTENT(out) :: x(:)
    DOUBLE PRECISION, INTENT(in)  :: b(:)
    DOUBLE PRECISION, INTENT(in)  :: omega
    INTEGER, OPTIONAL, INTENT(in) :: niter_in
!
    INTEGER          :: niter
    DOUBLE PRECISION :: rtolmx
!++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
!                        1.   No-preconditionning
!
    IF( omega .LT. 0.0d0 ) THEN
       x = b
       RETURN
    END IF
!++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
!                        2.   SSOR Preconditionner
!
    niter = 1
    rtolmx = 1.d-6
    IF(PRESENT(niter_in)) THEN
       niter = niter_in
    END IF
    CALL ssor(mat, b, omega, niter, rtolmx, x)
!++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
  END SUBROUTINE psolve
!
!===========================================================================
  SUBROUTINE ssor(mat, b, omega, nitmx, rtolmx, x, resid, nit, verbose)
!
!   Solve Ax = b using SSOR method
!
    USE cds
    TYPE(cds_mat)                           :: mat
    DOUBLE PRECISION, INTENT(out)           :: x(:)
    DOUBLE PRECISION, INTENT(in)            :: b(:)
    DOUBLE PRECISION, INTENT(in)            :: omega, rtolmx
    INTEGER, INTENT(in)                     :: nitmx
    DOUBLE PRECISION, OPTIONAL, INTENT(out) :: resid
    INTEGER, OPTIONAL, INTENT(out)          :: nit
    LOGICAL, OPTIONAL, INTENT(in)           :: verbose
!
    INTEGER          :: n, iter
    INTEGER          :: k, i, j, d, bw0, ny
    DOUBLE PRECISION :: omega1, bnrm2, residue
    DOUBLE PRECISION :: rhs(SIZE(x))
!++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
!                            1. Initialization
    n = SIZE(x)
    bw0 = SIZE(mat%rowv)
    ny = mat%ny
    bnrm2 = norm2(b)  ! Euclidian norm of RHS
    omega1 = 1.0d0-omega
    iter = 0
    DO
       iter = iter+1
!
!                            2. Forward SOR
!  Set RHS
       rhs = omega*b  ! rhs <- omega*b
       IF( iter .GT. 1 ) THEN
          rhs = rhs + omega1*x ! rhs <- rhs + (1-omega)*x0
          DO k=1,mat%ku        ! rhs <- rhs - omega*U*x0
             d = mat%dists(k)
             DO i=MAX(1,1-d),MIN(n,n-d)
                rhs(i) = rhs(i) - omega*mat%val(i,k)*x(i+d)
             END DO
          END DO
          IF( ny .NE. 0 ) THEN  ! Contributions from unicity BC
             rhs(ny) = rhs(ny) - omega*DOT_PRODUCT(mat%rowv(ny+1:bw0),x(ny+1:bw0))
          END IF
       END IF
!
!  Solve (1+omega*L) x = rhs
       x = rhs
       IF( ny .NE. 0 ) THEN  ! Contributions from unicity BC
          rhs(ny+1:bw0) = rhs(ny+1:bw0) - omega*mat%colh(ny:bw0)*x(ny)
       END IF
       DO i=ny+1,n
          DO k=-1,-mat%kl,-1
             d = mat%dists(k)
             j=i+d
             IF( j.LE.0 ) EXIT
             x(i) = x(i) - omega*mat%val(i,k)*x(j)
          END DO
       END DO
!
!                            3. Backward SOR
!  Set RHS
       rhs = omega*b + omega1*x  ! rhs <- omega*b + (1-omega)*x0
       IF( ny .NE. 0 ) THEN  ! Contributions from unicity BC
          rhs(ny+1:bw0) = rhs(ny+1:bw0) - omega*mat%colh(ny:bw0)*x(ny)
       END IF
       DO k=-mat%kl,-1           ! rhs <- rhs - omega*L*x0
          d = mat%dists(k)
          DO i=MAX(1,1-d),MIN(n,n-d)
             rhs(i) = rhs(i) - omega*mat%val(i,k)*x(i+d)
          END DO
       END DO
!
!  Solve (1+omega*U) x = rhs
       x = rhs
       DO i=n-1,ny+1,-1
          DO k=1,mat%ku
             d = mat%dists(k)
             j = i+d
             IF( j.GT.n ) EXIT
             x(i) = x(i) - omega*mat%val(i,k)*x(j)
          END DO
       END DO
       IF( ny .NE. 0 ) THEN  ! Contributions from unicity BC
          x(ny) = x(ny) - omega*DOT_PRODUCT(mat%rowv(ny+1:bw0),x(ny+1:bw0))
       END IF
!
!                            4. Compute residue
!
       IF( PRESENT(resid) ) THEN
          residue = norm2(b-vmx(mat,x)) / bnrm2
          IF(PRESENT(verbose)) THEN
             IF(verbose) WRITE(*,'(a,i8,1pe12.3)') 'it, resid', iter, residue
          END IF
          IF( residue .LT. rtolmx ) EXIT
       END IF
!
       IF( iter .GE. nitmx ) EXIT
    END DO  ! End of SSOR iterations
    IF(PRESENT(nit))  nit = iter
    IF(PRESENT(resid)) resid = residue
  END SUBROUTINE ssor
!===========================================================================
  SUBROUTINE diagbal(mat)
!
!   Diagonal matrix balancing: store D^(-1/2) in mat%bal
!
    USE cds
    TYPE(cds_mat)    :: mat
    INTEGER          :: n, bw0, ny, d, i, k
    DOUBLE PRECISION :: diag(mat%rank)
! 
    n = mat%rank
    ny = mat%ny
    bw0 = SIZE(mat%colh)
    diag(1:n) = mat%val(1:n,0)
    IF( MINVAL(diag) .LE. 0.0d0 ) THEN
       WRITE(*,'(a)') 'Diagonal elements of matrix are not stricly positive!'
       STOP
    END IF
    diag(1:n) = 1.0d0/SQRT(diag(1:n))
!
!     Scale the matrix

!$OMP parallel do private (k,d,i)
    DO k=-mat%kl,mat%ku
       d = mat%dists(k)
       DO i=MAX(1,1-d),MIN(n,n-d)
          mat%val(i,k) = diag(i)*diag(i+d)*mat%val(i,k)
       END DO
    END DO
!$OMP end parallel do
!
!     The ny^th column and row
    IF( ny.NE.0 ) THEN
       mat%rowv(1:bw0) = diag(1:bw0)*diag(ny)*mat%rowv(1:bw0)
       mat%colh(1:bw0) = diag(ny)*diag(1:bw0)*mat%colh(1:bw0)
    END IF
!
!     Save D^(-1/2)
    mat%bal(:) = diag(:)
  END SUBROUTINE diagbal
!++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
  FUNCTION norm2_1d(x)
!
!  Compute the 2-norm of 1d array
!
    DOUBLE PRECISION :: x(:)
    DOUBLE PRECISION :: norm2_1d
    norm2_1d = SQRT(SUM(x*x))
  END FUNCTION norm2_1d
!++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
  FUNCTION norm2_2d(x)
!
!  Compute the 2-norm of 2d array
!
    DOUBLE PRECISION :: x(:,:)
    DOUBLE PRECISION :: norm2_2d
    norm2_2d = SQRT(SUM(x*x))
  END FUNCTION norm2_2d
!++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
END MODULE tcdsmat_mod