!
! program to perform basic Linear Algebra routines
!
! R. Orvedahl 9.17-2014

!=========================================================================
! define precision type
!=========================================================================
module data_types
   implicit none
   integer, parameter, public :: dpt  = selected_real_kind(15,307)
   integer, parameter, public :: i8_t = selected_int_kind(15)
   integer, parameter, public :: spt  = selected_real_kind(6,37)
   integer, parameter, public :: i4_t = selected_int_kind(9)
end module data_types

!=========================================================================
! define variables that are used rather often
!=========================================================================
module variables
   use data_types
   implicit none
   real(kind=dpt), parameter :: zero = 0.0_dpt
   real(kind=dpt), parameter :: one  = 1.0_dpt
   real(kind=dpt), parameter :: two  = 2.0_dpt
   real(kind=dpt), parameter :: half = 2.0_dpt
end module variables

!=========================================================================
! main linear algebra module
!=========================================================================
module linear_algebra

   use data_types
   use variables

   implicit none

   private
   public :: solve_Ax_b, tridiagonal
   public :: err_size, err_singular
   public :: igauss, ilu, iqr

   ! error codes
   integer, parameter :: err_size     = -1
   integer, parameter :: err_singular = -2

   ! identifiers for what matrix method to use
   integer, parameter :: igauss = 1
   integer, parameter :: ilu    = 2
   integer, parameter :: iqr    = 3


   contains


   !-------------------------------------------------------------------
   ! wrapper to different elimination routines
   !-------------------------------------------------------------------
   subroutine solve_Ax_b(n, A, b, x, method, stat, det)

      ! INPUT
      integer,        intent(in) :: n, method
      real(kind=dpt), intent(in) :: A(n,n), b(n)

      ! OUTPUT
      real(kind=dpt), intent(out) :: x(n)
      integer, intent(out) :: stat
      real(kind=dpt), optional, intent(inout) :: det

      ! call apropriate routine
      if (method == igauss) then
          if (present(det)) then
             call gauss_elim(n, A, b, x, stat, det)
          else
             call gauss_elim(n, A, b, x, stat)
          endif

      elseif (method == ilu) then
         write(*,*)
         write(*,*) "ERROR: LU matrix solver not coded yet"
         write(*,*)
         stop
      elseif (method == iqr) then
         write(*,*)
         write(*,*) "ERROR: QR matrix solver not coded yet"
         write(*,*)
         stop
      else
         write(*,*)
         write(*,'(a,i3)') "ERROR: unknown matrix solver: ",method
         write(*,*)
         stop
      endif

      return

   end subroutine solve_Ax_b


   !-------------------------------------------------------------------
   ! solve a tridiagonal system of the form:
   !    a(i)*x(i-1) + b(i)*x(i) + c(i)*x(i+1) = rhs(i)
   !
   !    a(1) & c(n) are never used, so they are arbitrary
   !-------------------------------------------------------------------
   subroutine tridiagonal(n, a, b, c, rhs, phi)

      integer, intent(in) :: n
      real(kind=dpt), intent(in) :: a(n), b(n), c(n), rhs(n)
      real(kind=dpt), intent(inout) :: phi(n)

      real(kind=dpt) :: cprime(n), rhsprime(n), m
      integer :: i

      ! See en.wikipedia.org/wiki/Tridiagonal_matrix_algorithm
      !
      ! initialize first elements & avoid overwriting c, rhs
      rhsprime(1) = rhs(1)/b(1)
      cprime(1) = c(1)/b(1)

      ! set coefficients
      do i=2,n-1
         m = b(i) - a(i)*cprime(i-1)
         cprime(i) = c(i)/m
         rhsprime(i) = (rhs(i) - a(i)*rhsprime(i-1)) / m
      enddo
      rhsprime(n) = (rhs(n) - a(n)*rhsprime(n-1)) / (b(n) - a(n)*cprime(n-1))

      ! back sub
      phi(n) = rhsprime(n)
      do i=n-1,1,-1
         phi(i) = rhsprime(i) - cprime(i)*phi(i+1)
      enddo

      return

   end subroutine tridiagonal


   !-------------------------------------------------------------------
   ! perform gauss elimination
   !-------------------------------------------------------------------
   subroutine gauss_elim(n, Ain, b_in, x, stat, det)

      ! INPUT
      integer,        intent(in) :: n
      real(kind=dpt), intent(in) :: Ain(n,n), b_in(n)

      ! OUTPUT
      integer, intent(out) :: stat
      real(kind=dpt), intent(out) :: x(n)
      real(kind=dpt), optional, intent(inout) :: det

      ! LOCAL
      real(kind=dpt) :: A(n,n), b(n), scales(n), tmp
      real(kind=dpt) :: pivot, coeff, sum_soln
      integer :: i, j, k, l, nrowswap = 0

      ! A must be square
      if (size(Ain, 2) /= size(Ain, 1)) then
         write(*,*)
         write(*,*) "ERROR: Gauss-Elim => A must be square---"
         write(*,*)
         stat = err_size
         return
      endif

      ! b must be same size as A
      if (size(b_in) /= size(Ain, 1)) then
         write(*,*)
         write(*,*) "ERROR: Gauss-Elim => b must be same size as A---"
         write(*,*)
         stat = err_size
         return
      endif

      ! local copy of A
      A = Ain
      b = b_in

      ! main loop over rows
      do k=1,n-1

         ! find max value in each row
         do i=k,n
            scales(i) = zero
            do j=k,n
               scales(i) = max(scales(i), abs(A(i,j)))
            enddo
         enddo

         ! get row with largest pivoting element
         pivot = abs(A(k,k)/scales(k))
         l = k
         do j=k+1,n
            if (abs(A(j,k)/scales(j)) > pivot) then
               pivot = abs(A(j,k)/scales(j))
               l = j
            endif
         enddo

         ! check if system is singular
         if (pivot == zero) then
            write(*,*)
            write(*,*) "ERROR: Matrix is singular---"
            write(*,*)
            stat = err_singular
            return
         endif

         ! interchange rows k and l (if needed)
         if (l /= k) then
            ! this loop does the entire row-swap for A
            do j=k,n
               tmp = A(k,j)
               A(k,j) = A(l,j)
               A(l,j) = tmp
            enddo
            ! this swaps b
            tmp = b(k)
            b(k) = b(l)
            b(l) = tmp
            nrowswap = nrowswap + 1
         endif

         ! do forward elimination (after scaling and pivoting)
         do i=k+1,n

            coeff = A(i,k)/A(k,k)

            do j=k+1,n
               A(i,j) = A(i,j) - coeff*A(k,j)
            enddo

            A(i,k) = zero
            b(i) = b(i) - coeff*b(k)

         enddo

      enddo

      ! back substitution
      x(n) = b(n)/A(n,n)
      do i=n-1,1,-1

         sum_soln = b(i)
         do j=i+1,n
            sum_soln = sum_soln - A(i,j)*x(j)
         enddo

         x(i) = sum_soln/A(i,i)

      enddo

      ! determinant = (product of diagonals of A) * (-1)**nrowswaps
      if (present(det)) then
          det = one
          do i=1,n
             det = det*A(i,i)
          enddo

          det = det*(-one)**nrowswap
      endif

      return

   end subroutine gauss_elim

end module linear_algebra


!=========================================================================
! main program to test the above routines
!=========================================================================
program test_lin_alg

   use data_types
   use variables
   use linear_algebra

   implicit none

   ! local variables
   integer, parameter :: n = 40
   integer :: i, j, stat
   real(kind=dpt) :: A(n,n), x(n), rhs(n), xexact(n), rand(4), l2norm
   real(kind=dpt) :: acoef(n), bcoef(n), ccoef(n)

   call random_seed()

   ! initialize the array
   do i=1, n
      do j=1,n
         call random_number(rand) ! rand is an array of length 4
         A(i,j) = (-one)**(mod(i*j, 3))*(two*rand(1) - one)
      enddo
   enddo

   ! initialize x
   call random_number(xexact)
   xexact = 10.0_dpt*xexact - 5.0_dpt

   ! calculate RHS
   rhs = matmul(A, xexact)

   !-------------------------------------------------------
   ! test the A.x = b solver
   !-------------------------------------------------------
   call solve_Ax_b(n, A, rhs, x, igauss, stat)

   ! calculate L2 Norm: l2norm = sqrt( sum( (x(i) - xexact(i))**2 ) )
   l2norm = zero
   do i=1,n
      l2norm = l2norm + (x(i) - xexact(i))**2
   enddo
   l2norm = sqrt(l2norm)

   print *,""
   print *,"Test A.x = rhs systems"
   print *,""
   print *,"Size of A:",n
   print *,""
   print *,"Full Matrix:"
   print *,"   L2 Norm: ", l2norm

   !-------------------------------------------------------
   ! test the Tridiagonal solver
   !-------------------------------------------------------
   ! first setup a Tridiagonal system
   !   a(i)*x(i-1) + b(i)*x(i) + c(i)*x(i+1) = rhs(i)
   call random_number(acoef)
   call random_number(bcoef)
   call random_number(ccoef)
   ! make the system diagonally dominant for numerical stability
   !    |a(k)| >= |b(k)| + |c(k)|   k=2,n-1
   !    |a(1)| >  |c(1)|
   !    |a(n)| >  |b(n)|
   acoef = two*acoef - one
   bcoef = two*bcoef + two
   ccoef = two*ccoef - one

   ! fill interior of A
   A(:,:) = zero
   do i=2,n-1
      A(i,i-1) = acoef(i)
      A(i,i)   = bcoef(i)
      A(i,i+1) = ccoef(i)
   enddo
   ! fill corners
   A(1,1)   = bcoef(1)
   A(n,n)   = bcoef(2)
   A(1,2)   = ccoef(1)
   A(n,n-1) = acoef(n)

   ! get the RHS
   call random_number(xexact)
   xexact = 10.0_dpt*xexact - 5.0_dpt
   rhs = matmul(A,xexact)

   ! solve system
   call tridiagonal(n, acoef, bcoef, ccoef, rhs, x)

   ! calculate L2 Norm: l2norm = sqrt( sum( (x(i) - xexact(i))**2 ) )
   l2norm = zero
   do i=1,n
      l2norm = l2norm + (x(i) - xexact(i))**2
   enddo
   l2norm = sqrt(l2norm)

   print *,""
   print *,"Tridiagonal Matrix:"
   print *,"   L2 Norm: ", l2norm
   print *,""

end program test_lin_alg