! Program Example from "Essential Fortran 90 & 95" by Loren P. Meissner
! Copyright 1996. Copying for sale requires permission from the author.
! Otherwise, distribution is permitted if these three lines are included.

!  Example 8.5. Gauss elimination; partial pivoting with vector subscripts.
!  -- ENHANCED --

  module I05M

  ! specification part
    implicit none
    public :: Factor, For_Bak, Residual, Solve, Unsolve
    private :: Permute
    integer, parameter, public :: LOW = selected_real_kind( 6 ), HIGH = selected_real_kind( 12 )
    integer, save, private :: N
    real (kind = LOW), dimension(:, :),allocatable, save, private :: Save_A, LU
    integer, dimension(:), allocatable, save, private :: P
    logical, private, save :: Saved = .false., Exist = .false.

! procedure part
  contains

! >>************************************************************************************************
    subroutine Factor( A, B, X, E, X_Norm, E_Norm, Flag )   ! Module procedure (public)
      ! This subroutine solves the system of simultaneous linear algebraic equations
      !   A x = b, where A is an n by n matrix and x and b are n by m matrices.
      !   A and B are assumed-shape arrays supplied by the calling program.
      !   X, E, X_Norm, E_Norm, and Flag are computed and returned by Factor.
      !   E is the error computed as A e = r, where r = B - A x(calculated).
      !   Factor does not change A or B.
      ! On return from Factor, X contains the solution and E contains the error.
      !   X_Norm and E_Norm contain the (Euclidean vector) norms of X and E,
      !   respectively.  However, if Flag is false, the contents of X, E, X_Norm, and
      !   E_Norm are unpredictable, since the algorithm may have short-circuited
      !   somewhere in the middle.
      ! The problem size N, a copy of A, and the LU decomposition matrix (with
      !   permutation vector P) are saved in the module as private variables, and are
      !   not changed except when Factor is called.
      ! The subroutine Solve may be called to solve additional systems with the same
      !   matrix A, between successful (i.e., Flag = true) calls to Factor.
      !   Call Solve with a new right-hand side B whose first dimension equals the
      !   original N. Solve will compute the solution vector X, the corresponding
      !   error E, and the norms.
      ! DUMMY ARGUMENTS
      real (kind = LOW), dimension(:, :), intent (in) :: A, B
      real (kind = LOW), dimension(:, :), intent (out) :: X, E
      real (kind = LOW), dimension(:), intent (out) :: X_Norm, E_Norm
      logical, intent(out) :: Flag
      ! LOCAL DATA
      real (kind = LOW), dimension(size( A, dim = 1 )) :: S
      real (kind = LOW), dimension(size( B, dim = 2 )) :: R_Norm
      integer :: I, K  ! K is an ACID variable
      integer, dimension(1) :: Pivot
  ! start subroutine Factor

      ! Determine problem size; initialize private arrays
      N = size( A, dim = 1 )
      if ((size( A, dim = 2 ) /= N) .or. (size( B, dim = 1 ) /= N) .or. &
          any( shape( X ) /= shape( B ) ) .or. any( shape( E ) /= shape( B ) ) .or. &
          (size( X_Norm ) /= size( B, dim = 2 )) .or. (size( E_Norm ) /= size( B, dim = 2 )) )  then
        write (unit = *, fmt = *) " FACTOR size incompatibility."   ! Error
        stop
      end if
      allocate( Save_A(N, N), LU(N, N), P(N) )
      Exist = .true.
      Save_A = A
      LU = A
      P = (/ (K, K = 1, N) /)
      ! Store row norms in S
      S = maxval( abs( LU ), dim = 2 )

      if (any( S == 0 )) then                                                   ! Error
        write (unit = *, fmt = *) " Row norms: ", S
        write (unit = *, fmt = *) " One of the row norms is zero. "
        stop
      end if

      ! Perform factorization A = L * U with scaled partial pivoting.
      do I = 1, N
        ! Reduce column I
        LU(P(I: ), I) = real( real( LU(P(I: ), I), kind = HIGH ) &
          - matmul( real( LU(P(I: ), 1: I - 1), kind = HIGH ), real( LU(P(1: I - 1), I), kind = HIGH ) ), kind = LOW )
        if (all( abs( LU(P(I: ), I) ) <= 0.0 )) then
          write (unit = *, fmt = *) " All pivot candidates are zero. "
          stop
        end if
        Pivot = maxloc( abs( LU(P(I: ), I) ) / S(P(I: )) )
        call Permute( I, I - 1 + Pivot(1) )
        ! Reduce row I
        LU(P(I), I + 1: ) = real( real( LU(P(I), I + 1: ), kind = HIGH ) &
          - matmul( real( LU(P(I), 1: I - 1), kind = HIGH ), real( LU(P(1: I - 1), I + 1: ), kind = HIGH ) ), kind = LOW) &
          / LU(P(I), I)
      end do

      write (unit = *, fmt = *) " LU Factors"
      do I = 1, N
        write (unit = *, fmt = *) LU(I, 1: I - 1), "|", LU(I, I: )
      end do

      ! Apply LU with P to right-hand side B; compute residual (private) and error.
      Saved = .true.                                  ! Module has stored the Factors.
      call Solve( B, X, E, X_Norm, E_Norm, R_Norm )
      write (unit = *, fmt = *) " R_Norm [from module]: ", R_Norm
      Flag = all( E_Norm < X_Norm / 2.0 )
      Saved = Flag                ! Module has stored the Factors ready to Solve again
      return
    end subroutine Factor

! >>**********************************************************************************************
    subroutine Permute( I, J )          ! Module procedure (private)
      integer, intent(in) :: I, J
      integer :: Temp
  ! start subroutine Permute
      Temp = P(I)
      P(I) = P(J)
      P(J) = Temp
      return
    end subroutine Permute

! >>************************************************************************************************
    subroutine Unsolve ( )          ! Module procedure (public)
  ! start subroutine Unsolve
      Saved = .false.
      if (Exist) then
        deallocate( Save_A, LU, P )
      end if
      Exist = .false.
      return
    end subroutine Unsolve

! >>************************************************************************************************
    subroutine Solve( B, X, E, X_Norm, E_Norm, R_Norm )          ! Module procedure (public)
      ! DUMMY ARGUMENTS
      real (kind = LOW), dimension(:, :), intent (in) :: B
      real (kind = LOW), dimension(:, :), intent (out) :: X, E
      real (kind = LOW), dimension(:), intent (out) :: X_Norm, E_Norm, R_Norm
      !   LOCAL VARIABLES
      real (kind = LOW), dimension(size( B, dim = 1 ), size( B, dim = 2 )) :: R
  ! start subroutine Solve
      if ((.not. Saved) .or. (size( B, dim = 1 ) /= N) .or. &
          any( shape( X ) /= shape( B ) ) .or. any( shape( E ) /= shape( B ) ) .or. &
          (size( X_Norm ) /= size( B, dim = 2 )) .or. (size( E_Norm ) /= size( B, dim = 2 )) )  then
        write (unit = *, fmt = *) " SOLVE argument incompatibility."   ! Error
        X = 0.0_LOW
        E = 0.0_LOW
        X_Norm = 0.0_LOW
        E_Norm = huge( 1.0_LOW )
        return
      end if
      call For_Bak( B, X, X_Norm )
      call Residual( B, X, R, R_Norm )
      call For_Bak( R, E, E_Norm )
      return
    end subroutine Solve

! >>**********************************************************************************************
    subroutine For_Bak( B, X, X_Norm )          ! Module procedure (private)
      ! This subroutine performs forward and backward substitution,
      !   using LU and P stored at the latest call to Factor.
      ! B is M right-hand side vectors.
      ! The result will be returned in X, and its Euclidean norms in X_Norm.
      ! DUMMY ARGUMENTS
      real (kind = LOW), dimension(:, :), intent (in) :: B
      real (kind = LOW), dimension(:, :), intent (out) :: X
      real (kind = LOW), dimension(:), intent (out) :: X_Norm
      ! LOCAL DATA
      integer :: I
      real(kind = LOW), dimension(size( B, dim = 1), size( B, dim = 2)) :: C
  ! start subroutine For_Bak
      if ((.not. Saved) .or. (size( B, dim = 1 ) /= N) .or. &
          any( shape( X ) /= shape( B ) ) .or. (size( X_Norm ) /= size( B, dim = 2 )) )  then
        write (unit = *, fmt = *) " FOR_BAK argument incompatibility."   ! Error
        X = 0.0_LOW
        X_Norm = 0.0_LOW
        return
      end if
      do I = 1, N
        C(I, :) = real( real( B(P(I), :), kind = HIGH) &
          - matmul( real( LU(P(I), 1: I - 1), kind = HIGH ), real( C(1: I - 1, :), kind = HIGH) ), kind = LOW) / LU(P(I), I)
      end do
      do I = N, 1, -1
        X(I, :) = real( real( C(I, :), kind = HIGH ) &
          - matmul( real( LU(P(I), I + 1: ), kind = HIGH ), real( X(I + 1: , :), kind = HIGH ) ), kind = LOW )
      end do
      X_Norm = sqrt( sum( X ** 2, dim = 1 ) )
      return
    end subroutine For_Bak

! >>**********************************************************************************************
    subroutine Residual( B, X, R, R_Norm )          ! Module procedure (private)
      ! Uses Save_A from latest call to Factor
      ! DUMMY ARGUMENTS
      real(kind = LOW), intent (in), dimension(:, :) :: B, X
      real(kind = LOW), intent (out), dimension(:, :) :: R
      real(kind = LOW), intent (out), dimension(:) :: R_Norm
      ! LOCAL DATA
      integer :: I
 !  start subroutine Residual
      if ((.not. Saved) .or. (size( B, dim = 1 ) /= N) .or. &
          any( shape( X ) /= shape( B ) ) .or. any( shape( R ) /= shape( B ) ) .or. &
          (size( R_Norm ) /= size( B, dim = 2 )))  then
        write (unit = *, fmt = *) " RESIDUAL argument incompatibility."   ! Error
        R = 0.0_LOW
        R_Norm = 0.0_LOW
        return
      end if
      do I = 1, N
        R = real( real( B, kind = HIGH ) &
          - matmul( real( Save_A, kind = HIGH ), real( X, kind = HIGH ) ), kind = LOW )
      end do
      R_Norm = sqrt( sum( R ** 2, dim = 1 ) )
      return
    end subroutine Residual

  end module I05M

