module tridiag
  use set_precision, only: wp
  implicit none

  ! tdg data structure for tridiagonal matrices
  ! representation of n x n tridiagonal matrix A
  ! A(i,i) = a(i), A(i+1,i) = b(i), A(i,i+1) = c(i)
  ! array sizes: a(max_n), b(max_n-1), c(max_n-1)
  type :: tdg
     integer n, max_n
     real(kind=wp), pointer :: a(:), b(:), c(:)
  end type tdg
  
  
contains
  ! tridiagonal functions for tdg structure

  ! tdg_get(n) -- returns tdg structure for n x n tridiagonal matrix
  logical function tdg_get(mat,n)
    integer err, n
    type(tdg), intent(out) :: mat
  
    tdg_get = .true.
    allocate(mat%a(n), mat%b(n-1), mat%c(n-1), STAT=err)
    if ( err /= 0 ) then
       tdg_get = .false.
       return
    end if
    mat%n = n
    mat%max_n = n
  end function tdg_get
  !-----------------------------------------------
  

  ! tdg_free(mat) -- deallocates tdg structure mat
  subroutine tdg_free(mat)
    type(tdg):: mat
  
    deallocate(mat%a, mat%b, mat%c)
  end subroutine tdg_free
  !-----------------------------------------------
  

  ! tdg_resize(mat,new_n) -- re-allocates mat to be new_n x new_n
  !    -- no copying is done so old data is lost
  logical function tdg_resize(mat,new_n)
    integer :: err, new_n
    type(tdg) :: mat
  
    if ( new_n > mat%max_n ) then
       deallocate(mat%a, mat%b, mat%c)
       allocate(mat%a(new_n), mat%b(new_n-1), mat%c(new_n-1), &
            STAT = err)
       mat%max_n = new_n
    end if
    mat%n = new_n
  end function tdg_resize
  !-----------------------------------------------
  

  ! tdg_copy(mat,copy) -- copies mat to copy (both tdg structures)
  !    -- copy is resized to be mat%n x mat%n
  logical function tdg_copy(mat,copy)
    type(tdg) :: mat, copy
    integer i
  
    tdg_copy = .true.
    if ( .not. tdg_resize(copy,mat%n) ) then
       tdg_copy = .false.
       return
    end if
    do i = 1, mat%n - 1
       copy%a(i) = mat%a(i)
       copy%b(i) = mat%b(i)
       copy%c(i) = mat%c(i)
    end do
    copy%a(mat%n) = mat%a(mat%n)
    
  end function tdg_copy
  !-----------------------------------------------
  

  ! tdg_print(mat,output) -- prints tdg structure mat on unit output
  subroutine tdg_print(mat,output)
    integer i, output
    type(tdg) :: mat
  
    write (output,*) 'Tridiagonal matrix: n:', mat%n
    write (output,*) 'a:'
    write (output,'(3G22.15)') (mat%a(i), i = 1, mat%n)
    write (output,*) 'b:'
    write (output,'(3G22.15)') (mat%b(i), i = 1, mat%n-1)
    write (output,*) 'c:'
    write (output,'(3G22.15)') (mat%c(i), i = 1, mat%n-1)
  end subroutine tdg_print
  !-----------------------------------------------
  

  ! tdg_lu -- carries out LU factorization without pivoting on mat
  !   -- the L factor has ones on the diagonal implicitly
  !   -- mat is overwritten by the factors in compact form:
  !      L is in the strictly lower triangular part of mat and
  !      U is in the upper triangular part of mat
  !   -- Note: unless mat is diagonally dominant this might not
  !      give acceptable accuracy
  subroutine tdg_lu(mat,iflag)
    type(tdg) :: mat
    integer, intent(out) :: iflag
    integer :: i, n
  
    n = mat%n
    if ( mat%a(1) == 0 ) then
       iflag = -1
       return
    end if
    mat%b(1) = mat%b(1) / mat%a(1)
  
    do i = 2, n-1
       if ( mat%a(i) == 0 ) then
          iflag = -1
          return
       end if
       mat%a(i) = mat%a(i) - mat%b(i-1)*mat%c(i-1)
       mat%b(i) = mat%b(i) / mat%a(i)
    end do
    mat%a(n) = mat%a(n) - mat%b(n-1)*mat%c(n-1)
  end subroutine tdg_lu
  !-----------------------------------------------
  

  ! tdg_lusolve(mat,y,iflag)
  !   -- solves the system L*U*x = y (overwriting y with x)
  !                where mat contains L & U in compact form
  !   -- if mat%a(i) == 0 for some i, indicating U is singular,
  !      then set iflag to -1; otherwise iflag == 0 indicating success
  subroutine tdg_lusolve(mat,y,iflag)
    type(tdg), intent(in) :: mat
    real(kind=wp) :: y(:)
    integer, intent(out) :: iflag
    integer i
    logical has_zero
  
    iflag = 0
  
    ! Solve U.z = y, overwriting y with z
    do i = 2, mat%n
       y(i) = y(i) - mat%b(i-1)*y(i-1)
    end do
  
    ! check mat%a for zeros
    has_zero = .false.
    do i = 1, mat%n
       has_zero = has_zero .or. (mat%a(i) == 0)
    end do
    if ( has_zero ) then
       iflag = -1
       return
    end if
  
    ! Solve L.x = z (stored in y), overwriting y with x
    y(mat%n) = y(mat%n) / mat%a(mat%n)
    do i = mat%n-1, 1, -1
       y(i) = (y(i) - mat%c(i)*y(i+1)) / mat%a(i)
    end do
  end subroutine tdg_lusolve
  !-----------------------------------------------
  

  ! tdg_mvmult(mat,x,y) -- y <- mat*x
  subroutine tdg_mvmult(mat,x,y)
    type(tdg) :: mat
    real(kind=wp) x(:), y(:)
    integer i
  
    if ( mat%n == 0 ) return
    if ( mat%n == 1 ) then
       y(1) = mat%a(1)*x(1)
       return
    end if
    ! mat%n > 1
    y(1) = mat%a(1)*x(1) + mat%c(1)*x(2)
    do i = 2, mat%n - 1
       y(i) = mat%b(i-1)*x(i-1) + mat%a(i)*x(i) + mat%c(i)*x(i+1)
    end do
    y(mat%n) = mat%b(mat%n-1)*x(mat%n-1) + mat%a(mat%n)*x(mat%n)
  end subroutine tdg_mvmult
  !-----------------------------------------------
  

  ! tdg_lumult(mat,prod) -- prod <- L*U, 
  !   -- L & U stored in compact form in mat
  !      (L has an implicit unit diagonal)
  logical function tdg_lumult(mat,prod)
    type(tdg) :: mat, prod
    integer i
  
    tdg_lumult = .true.
    if ( .not. tdg_resize(prod,mat%n) ) then
       tdg_lumult = .false.
       return
    end if
  
    if ( mat%n <= 0 ) return
  
    prod%a(1) = mat%a(1)
    if ( mat%n == 1 ) return
    prod%c(1) = mat%c(1)
    do i = 2, mat%n - 1
       prod%b(i-1) = mat%b(i-1)*mat%a(i-1)
       prod%a(i)   = mat%a(i) + mat%b(i-1)*mat%c(i-1)
       prod%c(i)   = mat%c(i)
    end do
    prod%b(mat%n-1) = mat%b(mat%n-1)*mat%a(mat%n-1)
    prod%a(mat%n) = mat%a(mat%n) + &
         mat%b(mat%n-1)*mat%c(mat%n-1)
  end function tdg_lumult
  

  ! tdg_norm_diff -- returns max_{i,j}|mat1(i,j)-mat2(i,j)|
  !   -- this is *not* the 1-norm or the infinity norm
  function tdg_norm_diff(mat1,mat2) result (max_diff)
    type(tdg) :: mat1, mat2
    real(kind=wp) diff1, diff2, diff3, max_diff
    integer i, n
  
    if ( mat1%n /= mat2%n ) then
       max_diff = abs(mat1%n-mat2%n)
       return
    end if
  
    n = mat1%n
  
    max_diff = 0
    do i = 1, n-1
       diff1 = abs(mat1%a(i)-mat2%a(i))
       diff2 = abs(mat1%b(i)-mat2%b(i))
       diff3 = abs(mat1%c(i)-mat2%c(i))
       max_diff = max(max_diff,diff1,diff2,diff3)
    end do
    max_diff = max(max_diff,abs(mat1%a(n)-mat2%a(n)))
  end function tdg_norm_diff
  !-----------------------------------------------
  

  ! tridiag test routines
  subroutine tdg_test1
    integer, parameter :: n = 10
    real(kind=wp) y(n), x(n)
    type(tdg) :: mat1, mat2, temp
    integer i, iflag
    integer, parameter :: stdout = 6
    real(kind=wp) diff, max_diff
    real(kind=wp) sol(n), rhs(n), temp2(n)
  
    if ( .not. tdg_get(mat1,n) .or. .not. tdg_get(mat2,n) ) then
       print *, 'Cannot allocate tridiagonal matrix'
       return
    end if
    do i = 1, n-1
       mat1%a(i) = i+1
       mat1%b(i) = -0.5d0*i
       mat1%c(i) = -1.0d0/(1+i)
    end do
    mat1%a(n) = n+1
    ! print out the test matrix
    call tdg_print(mat1,stdout)
    ! save a copy
    if ( .not. tdg_copy(mat1,mat2) ) then
       print *, 'Error copying matrices'
    end if
    ! check that it's a copy
    call tdg_print(mat2,stdout)
    print *, 'Sizes of mat1 & mat2: ', mat1%n, mat2%n
    print *, 'Difference between copy and original =', &
         tdg_norm_diff(mat1,mat2)
    ! factor original matrix
    call tdg_lu(mat1,iflag)
    if ( iflag /= 0 ) then
       print *, 'Error in computing LU factorization'
    end if
    ! print out the factors in compact form
    print *, 'L, U factors in compact form'
    call tdg_print(mat1,stdout)
    ! multiply out factors, and compare with copy
    if ( .not. tdg_get(temp,mat1%n) ) then
       print *, 'Cannot allocate temp tridiagonal matrix'
       return
    end if
    if ( .not. tdg_lumult(mat1,temp) ) then
       print *, 'Error performing L.U multiplication'
    end if
    print *, 'L.U ='
    call tdg_print(temp,stdout)
    print *, 'Maximum difference between L.U & original =', &
         tdg_norm_diff(temp,mat2)
    
    ! set up linear system...
    rhs(1) = 1
    do i = 2, n
       rhs(i) = -2.0d0*rhs(i-1)/i
    end do
    ! and copy the right-hand side
    do i = 1, n
       sol(i) = rhs(i)
    end do
  
    ! solve linear system
    call tdg_lusolve(mat1,sol,iflag)
    if ( iflag /= 0 ) then
       print *, 'Apparently singular matrix'
    end if
    call tdg_mvmult(mat2,sol,temp2)
    max_diff = 0
    do i = 1, n
       max_diff = max(max_diff,abs(temp2(i)-rhs(i)))
    end do
    print *, 'Max norm error in reconstruction of rhs =', max_diff
  
    call tdg_free(mat1)
    call tdg_free(mat2)
    call tdg_free(temp)
  end subroutine tdg_test1
  !-----------------------------------------------
  
  

end module tridiag