! -*- mode: f90; coding: utf-8 -*-
! ----------------------------------------------------------------------------
! Copyright(c) 2002-2019 SPMODEL Development Group. All rights reserved.
!
! 履歴  2002/01/20  竹広真一
!       2002/06/10  竹広真一   ベクトル長問題回避のため lusol2 を使用
!       2005/01/10  竹広真一   msgdmp -> MessageNotify に変更
!       2006/03/04  竹広真一   コメントを RDoc 用に修正
!       2009/01/29  佐々木洋平 コメントを RDoc 用に修正
!       2009/08/06  竹広真一   ludecomp21 用ルーチン OMP 版変更
!       2013/12/14  佐々木洋平 interface 部を修正
!       2019/09/20  佐々木洋平 コメントを doxygen 用に修正, lapackと統合
! ----------------------------------------------------------------------------
!> @author Shin-ichi Takehiro, Youhei SASAKI
!> @copyright Copyright (C) SPMODEL Development Group, 2002-2026.
!> License is MIT/X11. see [COPYRIGHT](@ref COPYRIGHT) in detail
!>
!> @brief 行列の LU 分解による線形連立方程式の解法:
!> @details
!> spml/lumatrix モジュールは, LU 分解法により連立 1 次方程式を解くための
!> Fortran90 関数を提供する.
!>
!> 他のスペクトル計算用モジュールの中で登場する境界値問題を解くために用いられている.
!>
!> 同じ次数の複数個の連立 1 次方程式
!> \f[
!> A[ij] X [j] = B[i]
!> \f]
!> の解を, 同時に複数個の右辺ベクトル \f$B[i]\f$ に対して求めることが
!> できるようになっている.
!
module lumatrix
  use dc_types, only: DP
  use dc_message, only: MessageNotify
  implicit none
  private
  public LUDecomp
  public LUSolve
  !------------------------------------------------------------------------
  !> 与えられた行列の LU 分解を行い, ピボットを格納する.
  !>
  !> * LU 分解された結果は `Alu` に上書きされる.
  !>   そのピボット情報は `kp` に格納される.
  !>
  !> * LUSolve を用いる前にこのサブルーチンを呼んで `Alu` と `kp` を
  !>   計算しておく.
  interface LUDecomp
     module procedure ludecomp21
     module procedure ludecomp32
  end interface LUDecomp
  !-----------------------------------------------------------------------
  !> 連立 1 次方程式の解を求める
  !>
  !> @warning
  !> LUSolve を用いる前に LUDecompを呼んで Alu を LU 分解し,
  !> ピボット情報 kp を計算しておかねばならない.
  interface LUSolve
     module procedure lusolve211
     module procedure lusolve212
     module procedure lusolve322
     module procedure lusolve323
  end interface LUSolve
contains
  !-----------------------------------------------------------------------
  !> NDIM x NDIM の行列 ALU を LU 分解する.
  !> @note
  !> 結果は 入力行列に上書き
  subroutine ludecomp21(alu,kp)
    !> 入出力行列
    real(DP), intent(inout) :: alu(:,:)
    !> ピボット
    integer, intent(out)   :: kp(size(alu,1))
    if ( size(alu,1) > size(alu,2) ) then
      call MessageNotify('E','ludecomp',&
        'The third dimension is less than the second')
    elseif( size(alu,1) < size(alu,2) ) then
      call MessageNotify('W','ludecomp',&
        'The third dimension is grater than the second')
    endif
    call LUMAK1( alu, kp, size(alu,1) )
  end subroutine ludecomp21
  !-----------------------------------------------------------------------
  !> JDIM 個 の NDIM x NDIM の行列を一度に LU 分解.
  !> @note
  !> 結果は入力行列に上書き
  subroutine ludecomp32(alu,kp)
    !> 入力行列．JDIM が行列の個数, 後ろ二つの NDIMxNDIM が行列の形状
    real(DP), intent(inout) :: alu(:,:,:)
    !> ピボット
    integer, intent(out)    :: kp(size(alu,1),size(alu,2))
    if ( size(alu,2) > size(alu,3) ) then
      call MessageNotify('E','ludecomp',&
        'The third dimension is less than the second')
    elseif( size(alu,2) < size(alu,3) ) then
      call MessageNotify('W','ludecomp',&
        'The third dimension is grater than the second')
    endif
    call lumake( alu, kp, size(alu,1), size(alu,2) )
  end subroutine ludecomp32
  !-----------------------------------------------------------------------
  !> ALU(NDIM,NDIM), KP(NDIM), B(NDIM)
  !> NDIM x NDIM 型行列の連立方程式
  !> A X = B を 1 個の B に対して計算する.
  function lusolve211(alu,kp,b)
    !> 入力／LU 行列
    real(DP), intent(in)  :: alu(:,:)
    !> ピボット
    integer, intent(in)  :: kp(:)
    !> 右辺ベクトル
    real(DP), intent(in)  :: b(:)
    !> 解
    real(DP) :: lusolve211(size(b))
    lusolve211 = b
    call lusol2( lusolve211, alu , kp, 1, size(b) )
  end function lusolve211
  !-----------------------------------------------------------------------
  !> ALU(NDIM,NDIM), KP(NDIM), B(JDIM,NDIM)
  !> NDIM x NDIM 型行列の連立方程式
  !> A X = B を JDIM 個の B に対して計算する.
  function lusolve212(alu,kp,b)
    !> 入力／LU 行列
    real(DP), intent(in)  :: alu(:,:)
    !> ピボット
    integer,  intent(in)  :: kp(:)
    !> 右辺ベクトル
    real(DP), intent(in)  :: b(:,:)
    !> 解
    real(DP) :: lusolve212(size(b,1),size(b,2))
    lusolve212 = b
    call lusol2( lusolve212, alu , kp, size(b,1), size(b,2) )
  end function lusolve212
  !-----------------------------------------------------------------------
  !> ALU(JDIM,NDIM,NDIM), KP(JDIM,NDIM), B(JDIM,NDIM)
  !> NDIM x NDIM 型行列を JDIM 個並べた連立方程式
  !> A X = B をひとつの B の並びに対して計算する.
  function lusolve322(alu,kp,b)
    !> 入力／LU 行列
    real(DP), intent(in)  :: alu(:,:,:)
    !> ピボット
    integer,  intent(in)  :: kp(:,:)
    !> 右辺ベクトル
    real(DP), intent(in)  :: b(:,:)
    !> 解
    real(DP) :: lusolve322(size(b,1),size(b,2))
    lusolve322 = b
    call lusol3( lusolve322, alu, kp, size(b,1), size(b,2) )
  end function lusolve322
  !-----------------------------------------------------------------------
  !> ALU(JDIM,NDIM,NDIM), KP(JDIM,NDIM), B(IDIM,JDIM,NDIM)
  !> NDIM x NDIM 型行列を JDIM 個並べた連立方程式
  !> A X = B を IDIM 個の B に対して計算する.
  function lusolve323(alu,kp,b)
    !> 入力／LU 行列
    real(DP), intent(in)  :: alu(:,:,:)
    !> ピボット
    integer, intent(in)  :: kp(:,:)
    !> 右辺ベクトル
    real(DP), intent(in)  :: b(:,:,:)
    !> 解
    real(DP) :: lusolve323(size(b,1),size(b,2),size(b,3))
    lusolve323 = b
    call lusolv( lusolve323, alu , kp, size(b,1), size(b,2), size(b,3) )
  end function lusolve323
  ! -- LAPACK が無い場合 -------------------------------------------------------
  !---------------------------------------------------------------------
  !> LU 分解: ndim x ndim の行列 jdim 個を一度に計算
  !> @note
  !> LU 行列は入力に上書きされる
  subroutine lumake(alu, kp, jdim, ndim)
    implicit none
    integer, intent(in)    :: jdim
    integer, intent(in)    :: ndim
    integer, intent(out)   :: kp(jdim, ndim)
    real(DP), intent(inout) :: alu(jdim, ndim, ndim)
    integer :: j, k, m, n
    real(DP) :: pivot, temp
   !$omp parallel do default(shared) private(j,k,m,n,pivot,temp) schedule(static)
    do j = 1, jdim
      do k = 1, ndim-1
        pivot   = alu(j,k,k)
        kp(j,k) = k
        do m = k+1, ndim
          if (abs(alu(j,m,k)) > abs(pivot)) then
            pivot   = alu(j,m,k)
            kp(j,k) = m
          end if
        end do
        if (kp(j,k) /= k) then
          do n = 1, ndim
            temp             = alu(j,k,n)
            alu(j,k,n)       = alu(j,kp(j,k),n)
            alu(j,kp(j,k),n) = temp
          end do
        end if
        do n = k+1, ndim
          alu(j,k,n) = alu(j,k,n) / pivot
          do m = k+1, ndim
            alu(j,m,n) = alu(j,m,n) - alu(j,m,k) * alu(j,k,n)
          end do
        end do
      end do
    end do
    !$omp end parallel do
  end subroutine lumake
  !---------------------------------------------------------------------
  !> LU 分解: ndim x ndim の行列一個を計算
  !> @note
  !> LU 行列は入力に上書き
  subroutine lumak1(alu, kp, ndim)
    implicit none
    integer, intent(in)       :: ndim
    real(DP), intent(inout)   :: alu(ndim,ndim)
    integer, intent(out)      :: kp(ndim)
    integer  :: k, m, n
    real(DP) :: pivot, temp
    do k = 1, ndim-1
      pivot = alu(k,k)
      kp(k) = k
      do m = k+1, ndim
        if (abs(alu(m,k)) > abs(pivot)) then
          pivot = alu(m,k)
          kp(k) = m
        end if
      end do
      if (kp(k) /= k) then
        do n = 1, ndim
          temp      = alu(k,n)
          alu(k,n)  = alu(kp(k),n)
          alu(kp(k),n) = temp
        end do
      end if
      !$omp parallel default(shared)
      !$omp do
      do n = k+1, ndim
        alu(k,n) = alu(k,n) / pivot
      end do
      !$omp end do
      !$omp do
      do n = k+1, ndim
        do m = k+1, ndim
          alu(m,n) = alu(m,n) - alu(m,k) * alu(k,n)
        end do
      end do
      !$omp end do
      !$omp end parallel
    end do
  end subroutine lumak1
  !---------------------------------------------------------------------
  !> 求解: ndim x ndim 行列を jdim 個並べた連立方程式
  !> AX = B を idim 個 の Bについて計算する.
  !> @note
  !> 解は右辺入力ベクトルに上書き
  subroutine lusolv(xv, alu, kp, idim, jdim, ndim)
    implicit none
    integer, intent(in)     :: idim
    integer, intent(in)     :: jdim
    integer, intent(in)     :: ndim
    real(DP), intent(inout) :: xv(idim, jdim, ndim)
    real(DP), intent(in)    :: alu(jdim, ndim, ndim)
    integer, intent(in)     :: kp(jdim, ndim)
    integer                 :: i, j, k, n, nn
    real(DP)                :: temp
    !$omp parallel do default(shared) private(i,j,k,n,nn,temp) schedule(static)
    do j = 1, jdim
      do k = 1, ndim-1
        do i = 1, idim
          if (kp(j,k) /= k) then
            temp            = xv(i,j,k)
            xv(i,j,k)       = xv(i,j,kp(j,k))
            xv(i,j,kp(j,k)) = temp
          end if
        end do
      end do
      do n = 1, ndim
        do i = 1, idim
          xv(i,j,n) = xv(i,j,n) / alu(j,n,n)
          do nn = n+1, ndim
            xv(i,j,nn) = xv(i,j,nn) - xv(i,j,n) * alu(j,nn,n)
          end do
        end do
      end do
      do k = ndim-1, 1, -1
        do n = k+1, ndim
          do i = 1, idim
            xv(i,j,k) = xv(i,j,k) - xv(i,j,n) * alu(j,k,n)
          end do
        end do
      end do
    end do
    !$omp end parallel do
  end subroutine lusolv
  !---------------------------------------------------------------------
  !> 求解: ndim x ndim 行列を jdim 個並べた連立方程式
  !> AX = B を jdim 個 の Bについて計算する.
  !> @note
  !> 解は右辺入力ベクトルに上書き
  subroutine lusol3(xv, alu, kp, jdim, ndim)
    implicit none
    integer, intent(in)     :: jdim
    integer, intent(in)     :: ndim
    real(DP), intent(inout) :: xv(jdim, ndim)
    real(DP), intent(in)    :: alu(jdim, ndim, ndim)
    integer, intent(in)     :: kp(jdim, ndim)
    real(DP)                :: temp
    integer                 :: j, k, n, nn
    !$omp parallel do default(shared) private(j,k,n,nn,temp) schedule(static)
    do j = 1, jdim
      do k = 1, ndim-1
        if (kp(j,k) /= k) then
          temp         = xv(j,k)
          xv(j,k)      = xv(j,kp(j,k))
          xv(j,kp(j,k)) = temp
        end if
      end do
      do n = 1, ndim
        xv(j,n) = xv(j,n) / alu(j,n,n)
        do nn = n+1, ndim
          xv(j,nn) = xv(j,nn) - xv(j,n) * alu(j,nn,n)
        end do
      end do
      do k = ndim-1, 1, -1
        do n = k+1, ndim
          xv(j,k) = xv(j,k) - xv(j,n) * alu(j,k,n)
        end do
      end do
    end do
    !$omp end parallel do
  end subroutine lusol3
  !---------------------------------------------------------------------
  !> 求解: ndim x ndim 型行列の連立方程式 A X = B を idim 個の B に対して計算.
  !> [LUSOLV](@ref lusolv) の JDIM = 1 に相当
  !> @note
  !> * 解は右辺の入力ベクトルに上書きされる.
  !> * JDIM=1 には際のベクトル長が短くなるためこのルーチンを用意
  subroutine lusol2(xv, alu, kp, idim, ndim)
    implicit none
    integer, intent(in)      :: idim
    integer, intent(in)      :: ndim
    real(DP), intent(inout)  :: xv(idim, ndim)
    real(DP), intent(in)     :: alu(ndim,ndim)
    integer, intent(in)      :: kp(ndim)
    integer  :: i, k, n, nn
    real(DP) :: temp
    !$omp parallel do default(shared) private(i,k,temp) schedule(static)
    do i = 1, idim
      do k = 1, ndim-1
        if (kp(k) /= k) then
          temp       = xv(i,k)
          xv(i,k)    = xv(i,kp(k))
          xv(i,kp(k)) = temp
        end if
      end do
    end do
    !$omp end parallel do
    !$omp parallel do default(shared) private(i,n,nn) schedule(static)
    do i = 1, idim
      do n = 1, ndim
        xv(i,n) = xv(i,n) / alu(n,n)
        do nn = n+1, ndim
          xv(i,nn) = xv(i,nn) - xv(i,n) * alu(nn,n)
        end do
      end do
    end do
    !$omp end parallel do
    !$omp parallel do default(shared) private(i,k,n) schedule(static)
    do i = 1, idim
      do k = ndim-1, 1, -1
        do n = k+1, ndim
          xv(i,k) = xv(i,k) - xv(i,n) * alu(k,n)
        end do
      end do
    end do
    !$omp end parallel do
  end subroutine lusol2
end module lumatrix
