! ----------------------------------------------------------------------
! Testing LU decomposition by NR and LAPACK
! LAPACK sgesv/dgesv, sgetrf & getrs
! Numerical Recipes ludcmp & lubksb
! ----------------------------------------------------------------------
! Command-line compilation:
! Linux:
! gfortran -O3 mNR.for testLU.f90 -llapack (default LAPACK library)
! ifort -heap-arrays mNR.for testLU.f90 -mkl (Intel MKL library)
! nvfortran -O3 mNR.for testLU.f90 -lblas (Nvidia HPC SDK library)
! ----------------------------------------------------------------------
! Windows:
! gfortran -O3 mNR.for testLU.f90 -llibblas (PGI libraries .. a problem with single-precision BLAS)
! ifort -heap-arrays -o a mNR.for testLU.f90 -Qmkl (Intel MKL library)
! (call setvars.bat or ifortvars.bat)
! gfortran -O3 testLU.f90 "%MKLROOT%"\lib\intel64\mkl_rt.lib (Intel MKL library)
! ----------------------------------------------------------------------

! ----------------------------------------------------------------------
! Global constants
! ----------------------------------------------------------------------
MODULE mConst

! single and double precision, physical size of arrays
integer,parameter :: sp=4,dp=8,nmax=10000

logical,parameter,private :: T=.TRUE.,F=.FALSE.
logical,parameter :: runLAPACK=T
logical,parameter :: runNR=T

END MODULE

! ----------------------------------------------------------------------
! Wallclock timer
! Usage: call tic(); call task(); print *,toc()
! ----------------------------------------------------------------------
MODULE mTicToc

implicit none
private
public tic,toc
integer,parameter :: wp=8
real(wp) :: time

CONTAINS

subroutine tic()
integer count,count_rate
call system_clock(count,count_rate)
time=real(count,wp)/real(count_rate,wp)
end subroutine

real(wp) function toc()
integer count,count_rate
call system_clock(count,count_rate)
toc=real(count,wp)/real(count_rate,wp)-time
end function

END MODULE

! ----------------------------------------------------------------------
! LU decomposition via the LAPACK library
! Overloaded defined operator .LAPACKLU.
! Usage: c4=a4.LAPACKLU.b4; c8=a8.LAPACKLU.b8
!        a4 real(4) matrix, b4, c4 real(4) vectors
!        a8 real(8) matrix, b8, c8 real(8) vectors
! ----------------------------------------------------------------------
MODULE mLAPACK_LU

implicit none
private
public operator(.LAPACKLU.)
integer,parameter :: sp=4,dp=8
integer,parameter :: LAPACK_sv=1,LAPACK_trfs=2,LAPACK_svx=3
integer,parameter :: switch=LAPACK_sv

interface operator (.LAPACKLU.)
module procedure LAPACK_sp_LU,LAPACK_dp_LU
end interface

CONTAINS

function LAPACK_sp_LU(A,b) RESULT (x)
real(sp),intent(in) :: A(:,:),b(:)
real(sp) :: x(size(A,1))
real(sp),allocatable :: aa(:,:)
integer,allocatable :: ipiv(:)
integer :: lda,ldb,n,nrhs,info
character(1) :: trans

lda=size(A,1)
ldb=size(b,1)
n=lda
nrhs=1
trans='N' ! trans: 'N' for no op, 'T' for transpose, 'C' for conjugate transpose
x=b
allocate (aa(size(A,1),size(A,2))); aa=A  ! allocate (aa,source=A)
allocate (ipiv(n))
select case (switch)
case (LAPACK_sv)
  call sgesv(n,nrhs,aa,lda,ipiv,x,ldb,info)
  if (info/=0) print *,'gesv info: ',info
case (LAPACK_trfs)
  call sgetrf(n,n,aa,lda,ipiv,info)
  if (info/=0) print *,'getrf info: ',info
  call sgetrs(trans,n,nrhs,aa,lda,ipiv,x,ldb,info)
  if (info/=0) print *,'getrs info: ',info
end select
deallocate (aa,ipiv)
end function

function LAPACK_dp_LU(A,b) RESULT (x)
real(dp),intent(in) :: A(:,:),b(:)
real(dp) :: x(size(A,1))
real(dp),allocatable :: aa(:,:)
integer,allocatable :: ipiv(:)
integer :: lda,ldb,n,nrhs,info
character(1) :: trans

lda=size(A,1)
ldb=size(b,1)
n=lda
nrhs=1
trans='N' ! trans: 'N' for no op, 'T' for transpose, 'C' for conjugate transpose
x=b
allocate (aa(size(A,1),size(A,2))); aa=A  ! allocate (aa,source=A)
allocate (ipiv(n))
select case (switch)
case (LAPACK_sv)
  call dgesv(n,nrhs,aa,lda,ipiv,x,ldb,info)
  if (info/=0) print *,'gesv info: ',info
case (LAPACK_trfs)
  call dgetrf(n,n,aa,lda,ipiv,info)
  if (info/=0) print *,'getrf info: ',info
  call dgetrs(trans,n,nrhs,aa,lda,ipiv,x,ldb,info)
  if (info/=0) print *,'getrs info: ',info
end select
deallocate (aa,ipiv)
end function

END MODULE

! ----------------------------------------------------------------------
! LU decomposition via Numerical Recipes
! Defined operator .NRLU.
! Usage: c8=a8.NRLU.b8
!        a8 real(8) matrix, b8 real(8) vector
! ----------------------------------------------------------------------
MODULE mNR_LU
use mNR
implicit none
private
public operator(.NRLU.)
integer,parameter :: sp=4,dp=8
integer,parameter :: NR_LU=1,NR_GAUSSJ=2
integer,parameter :: switch=NR_LU

interface operator (.NRLU.)
module procedure NR_dp_LU
end interface

CONTAINS

function NR_dp_LU(A,b) RESULT (x)
real(dp),intent(in) :: A(:,:),b(:)
real(dp) :: x(size(b,1))
real(dp),allocatable :: aa(:,:)
integer,allocatable :: indx(:)
integer :: np,n,mp,m
real(dp) :: d

np=size(A,1)
n=np
mp=1
m=mp
x=b
allocate (aa(size(A,1),size(A,2))); aa=A  ! allocate (aa,source=A)
select case (switch)
case (NR_LU)
  allocate (indx(n))
  call ludcmp(aa,n,np,indx,d)
  call lubksb(aa,n,np,indx,x)
  deallocate (indx)
case (NR_GAUSSJ)
  call gaussj(aa,n,np,x,m,mp)
end select
deallocate (aa)
end function

END MODULE

! ----------------------------------------------------------------------
  PROGRAM TestLU
! ----------------------------------------------------------------------
use mConst
use mTicToc
use mLAPACK_LU
use mNR_LU
implicit none
real(dp),allocatable :: A(:,:),b(:),x(:)

ALLOCATE (A(nmax,nmax),b(nmax),x(nmax))
call random_number(A)

if (runLAPACK) then
  print *,'LAPACK'
  b=1.
  call tic()
  x=real(A,sp).LAPACKLU.real(b,sp)
  print *,'dt (SP) = ',toc()
  print '(1x,a,10e20.12)','x =',x(1:3)
  b=matmul(A,x)
  print '(1x,a,10e20.12)','b =',b(1:3)

  b=1.
  call tic()
  x=A.LAPACKLU.b
  print *,'dt (DP) = ',toc()
  print '(1x,a,10e20.12)','x =',x(1:3)
  b=matmul(A,x)
  print '(1x,a,10e20.12)','b =',b(1:3)
endif

if (runNR) then
  print *,'NR'
  b=1.
  call tic()
  x=A.NRLU.b
  print *,'dt (DP) = ',toc()
  print '(1x,a,10e20.12)','x =',x(1:3)
  b=matmul(A,x)
  print '(1x,a,10e20.12)','b =',b(1:3)
endif

END PROGRAM