module LINEAR_SOLVE
implicit none

public :: KROUT

integer, parameter, public  :: DP = selected_real_kind(12, 60)

contains

subroutine KROUT(MO, N, M, A, KA, B, KB, IERR)
!-----------------------------------------------------------------------
!  CROUT PROCEDURE FOR INVERTING MATRICES AND SOLVING EQUATIONS
!-----------------------------------------------------------------------
!  A IS A MATRIX OF ORDER N WHERE N IS GREATER THAN OR EQUAL TO 1.
!  IF MO = 0 THEN THE INVERSE OF A IS COMPUTED AND STORED IN A.
!  IF MO IS NOT 0 THEN THE INVERSE IS NOT COMPUTED.

!  IF M IS GREATER THAN 0 THEN B IS A MATRIX HAVING N ROWS AND M COLUMNS.
!  IN THIS CASE AX = B IS SOLVED AND THE SOLUTION X IS STORED IN B.
!  IF M=0 THEN THERE ARE NO EQUATIONS TO BE SOLVED.
!  N.B. B is passed as a VECTOR not as a matrix.

!  KA = THE LENGTH OF THE COLUMNS OF THE ARRAY A
!  KB = THE LENGTH OF THE COLUMNS OF THE ARRAY B (IF M > 0)

!  IERR IS A VARIABLE THAT REPORTS THE STATUS OF THE RESULTS. WHEN
!  THE ROUTINE TERMINATES IERR HAS ONE OF THE FOLLOWING VALUES ...
!     IERR =  0   THE REQUESTED TASK WAS PERFORMED.
!     IERR = -1   EITHER N, KA, OR KB IS INCORRECT.
!     IERR =  K   THE K-TH PIVOT ELEMENT IS 0.

!-----------------------------------------------------------------------

! Adapted from the routine KROUT in the NSWC Math. Library by Alan Miller
! Latest revision - 3 August 1998

integer, intent(in)                            :: MO
integer, intent(in)                            :: N
integer, intent(in)                            :: M
real (kind=DP), intent(in out), dimension(:,:) :: A     ! a(ka,n)
integer, intent(in)                            :: KA
real (kind=DP), intent(in out), dimension(:)   :: B
integer, intent(in)                            :: KB
integer, intent(out)                           :: IERR

!     INDEX IS AN ARRAY OF DIMENSION N-1 OR LARGER THAT IS USED BY THE
!     ROUTINE FOR KEEPING TRACK OF THE ROW INTERCHANGES THAT ARE MADE.
!     IF MO IS NOT 0 THEN THIS ARRAY IS NOT NEEDED.

!     TEMP IS AN ARRAY OF DIMENSION N OR LARGER THAT IS USED WHEN A
!     IS INVERTED.  IF MO IS NOT 0 THEN THIS ARRAY IS NOT NEEDED.

integer, allocatable, dimension(:)         :: INDX
real (kind=DP), allocatable, dimension(:)  :: TEMP

integer        :: I, J, JP1, K, KJ, KM1, KP1, L, LJ, MAXB, NJ, NMJ, NMK, NM1, &
                  ONEJ
real (kind=DP) :: D, DSUM, P, T
real (kind=DP), parameter :: ZERO = 0.0_DP, ONE = 1.0_DP

if (N < 1 .or. KA < N) then
  IERR = -1
  return
end if
if (M > 0 .and. KB < N) then
  IERR = -1
  return
end if

IERR = 0
if (N < 2) then

!                      CASE WHEN N = 1

  D = A(1,1)
  if (D == ZERO) then
    IERR = N
    return
  end if
  if (MO == 0) then
    A(1,1) = ONE / D
  end if

  if (M <= 0) then
    return
  end if
  MAXB = KB*M
  do KJ = 1,MAXB,KB
    B(KJ) = B(KJ)/D
  end do
  return
end if

!                      General case

if (MO == 0) then
  allocate( INDX(N-1), TEMP(N) )
end if

NM1 = N - 1
do K = 1,NM1
  KP1 = K + 1

!               SEARCH FOR THE K-TH PIVOT ELEMENT

  P = abs(A(K,K))
  L = K
  do I = KP1,N
    T = abs(A(I,K))
    if (P >= T) then
      cycle
    end if
    P = T
    L = I
  end do
  if (P == ZERO) then
!                  K-TH PIVOT ELEMENT IS 0
    IERR = K
    return
  end if

  P = A(L,K)
  if (MO == 0) then
    INDX(K) = L
  end if
  if (K /= L) then

!                  INTERCHANGING ROWS K AND L

    do J = 1,N
      T = A(K,J)
      A(K,J) = A(L,J)
      A(L,J) = T
    end do

    if (M > 0) then
      KJ = K
      LJ = L
      do J = 1,M
        T = B(KJ)
        B(KJ) = B(LJ)
        B(LJ) = T
        KJ = KJ + KB
        LJ = LJ + KB
      end do
    end if
  end if

!                  COMPUTE THE K-TH ROW OF U

  if (K <= 1) then
    do J = KP1,N
      A(K,J) = A(K,J)/P
    end do

  else
    do J = KP1,N
      DSUM = A(K,J) - dot_product( A(K,1:KM1), A(1:KM1,J) )
      A(K,J) = DSUM / P
    end do
  end if

!               COMPUTE THE (K+1)-ST COLUMN OF L

  do I = KP1,N
    DSUM = A(I,KP1) - dot_product( A(I,1:K), A(1:K,KP1) )
    A(I,KP1) = DSUM
  end do

  KM1 = K
end do

!                 CHECK THE N-TH PIVOT ELEMENT

if (A(N,N) == ZERO) then
  IERR = N
  return
end if

!                 SOLVING THE EQUATION LY = B

if (M > 0) then
  MAXB = KB*M
  do ONEJ = 1,MAXB,KB
    KJ = ONEJ
    B(KJ) = B(KJ)/A(1,1)
    do K = 2,N
      KJ = KJ + 1
      DSUM = B(KJ)
      KM1 = K - 1
      LJ = ONEJ
      do L = 1,KM1
        DSUM = DSUM - A(K,L)*B(LJ)
        LJ = LJ + 1
      end do
      B(KJ) = DSUM / A(K,K)
    end do
  end do

!                 SOLVING THE EQUATION UX = Y

  do NJ = N,MAXB,KB
    KJ = NJ
    do NMK = 1,NM1
      K = N - NMK
      LJ = KJ
      KJ = KJ - 1
      DSUM = B(KJ)
      KP1 = K + 1
      do L = KP1,N
        DSUM = DSUM - A(K,L)*B(LJ)
        LJ = LJ + 1
      end do
      B(KJ) = DSUM
    end do
  end do
end if

!               REPLACE L WITH THE INVERSE OF L

if (MO /= 0) then
  return
end if

do J = 1,NM1
  A(J,J) = ONE / A(J,J)
  JP1 = J + 1
  do I = JP1,N
    DSUM = dot_product( A(I,J:I-1), A(J:I-1,J) )
    A(I,J) = -DSUM / A(I,I)
  end do
end do
A(N,N) = ONE / A(N,N)

!           SOLVE UX = Y WHERE Y IS THE INVERSE OF L

do NMK = 1,NM1
  K = N - NMK
  KP1 = K + 1
  do J = KP1,N
    TEMP(J) = A(K,J)
    A(K,J) = ZERO
  end do

  do J = 1,N
    DSUM = A(K,J) - dot_product( TEMP(KP1:N), A(KP1:N,J) )
    A(K,J) = DSUM
  end do
end do

!                    COLUMN INTERCHANGES

do NMJ = 1,NM1
  J = N - NMJ
  K = INDX(J)
  if (J == K) then
    cycle
  end if
  do I = 1,N
    T = A(I,J)
    A(I,J) = A(I,K)
    A(I,K) = T
  end do
end do

if (MO == 0) then
  deallocate( INDX, TEMP )
end if

return
end subroutine KROUT

end module LINEAR_SOLVE



program TEST_KROUT
use LINEAR_SOLVE
implicit none

integer                     :: I, IERR, J, M, N, POS
integer, dimension(8)       :: TSTART, TEND
real (kind=DP)              :: SUMSQ
real                        :: ELAPSED
real (kind=DP), allocatable, dimension(:,:) :: A, AA
real (kind=DP), allocatable, dimension(:)   :: B

write(unit=*, fmt="(a)", advance="NO") " Enter size of matrix A: "
read(unit=*, fmt=*) N
write(unit=*, fmt="(a)", advance="NO") " How many RHS vectors?: "
read(unit=*, fmt=*) M

allocate( A(N,N), B(N*M), AA(N,N) )

! Fill A with random numbers.
! Set B so that the solution for column k is a vector of k's.

call date_and_time(values=TSTART)

call random_number( A )
AA = A
POS = 1
do J = 1, M
  do I = 1, N
    B(POS) = J * sum( A(I,:) )
    POS = POS + 1
  end do
end do

call KROUT(0, N, M, A, N, B, N, IERR)
write(unit=*, fmt=*) "ierr =", IERR
write(unit=*, fmt=*) "Solutions: "
POS = 1
do J = 1, M
  write(unit=*, fmt="(' ', 10f7.3)") B(POS:POS+N-1)
  POS = POS + N
end do

! Find the product of A & its inverse

A = matmul(A, AA)
SUMSQ = 0.0_DP
do I = 1, N
  do J = 1, N
    if (I ==J) then
      SUMSQ = SUMSQ + (1.0_DP - A(I,I))**2
    else
      SUMSQ = SUMSQ + A(I,J)**2
    end if
  end do
end do
write(unit=*, fmt="(a, g13.5)") " RMS error = ", sqrt(SUMSQ / N**2)

call date_and_time(values=TEND)
ELAPSED = 3600.0*(TEND(5) - TSTART(5)) + 60.0*(TEND(6) - TSTART(6)) + &
          (TEND(7) - TSTART(7)) + 0.001*(TEND(8) - TSTART(8))
write(unit=*, fmt="(a, f7.2, a)") " Elapsed time = ", ELAPSED, " sec."

stop
end program TEST_KROUT