MODULE Global_data
  !Symbolic names for kind types of single- and double-precision reals: 
  INTEGER, PARAMETER :: SP = KIND(1.0) 
  INTEGER, PARAMETER :: DP = KIND(1.0D0) 
  !Frequently used mathematical constants (with precision to spare): 
  REAL(DP), PARAMETER :: Pi=3.141592653589793238462643383279502884197_dp 
  ! parameters for the equation
  REAL(DP) :: x_min,x_max,h,kappa
END MODULE Global_data

MODULE finite_difference
! here we just need dp, kappa, h, and x_min 
! if we set wp=>sp we can change to single precision numbers
  USE global_data, ONLY : wp=>dp,h,kappa,x_min
  IMPLICIT NONE

CONTAINS

  !#################################################
  !    Solving matrix equation A x = b
  !    
  !    | b_0 c_0               | | x_0 |    | b_0 |
  !    | a_1 b_1 c_1           | | x_1 |    | b_1 |
  !    |     a_2 b_2 c_2       | | x_2 |    | b_2 |
  !    |       \  \   \        | |     |    |     |
  !    |     \   \   \         | |     |    |     |
  !    |    a_i b_i c_i        | | x_i |    | b_i |
  !    |     \   \   \         | |     |    |     |
  !    |        \   \   \      | |     |    |     |
  !    |           \    \  \   | |     |    |     |
  !    |               a_n b_n | | x_n |    | b_n |
  ! 
  !#################################################


  SUBROUTINE setup_lhs_coefficients(a,b,c,x,n)
    !    we assign y to have 'no_of_equations's elements throughout to keep the program consistent
    !   The ODE here is y'' + kappa y' + x y = 0
    !   so that the coefficients are
    !    alpha_i = 1/dx^2 - 0.5*kappa/dx
    !    beta_i = -2/dx^2 + x
    !    gamma_i = 1/dx^2 + 0.5*kappa/dx
    !

    INTEGER :: i,n
    REAL(wp) :: a(0:n),b(0:n),c(0:n),x(0:n)

    a(0) = 0.
    b(0) = 1.
    c(0) = 0.

    DO i = 1,n-1
       a(i) = 1._wp/h/h - 0.5_wp*kappa/h
       b(i) = -2._wp/h/h + x(i)  
       c(i) = 1._wp/h/h + 0.5_wp*kappa/h
    ENDDO

    a(n) = 0.
    b(n) = 1.
    c(n) = 0.
    
  END SUBROUTINE setup_lhs_coefficients

  SUBROUTINE setup_rhs_coefficients(d,x,n)
    ! Setup coefficients for the right hand side
    ! Here d_0 and d_n represent the boundary conditions at either end
 
    INTEGER :: i,n
    REAL(wp) :: d(0:n),x(0:n)
    
    !  here y(x_min) = 0
    d(0) = 0._wp
    DO i = 1,n-1
       d(i) = 0._wp
    ENDDO
    !  here y(x_max) = 1
    d(n) = 1._wp

  END SUBROUTINE setup_rhs_coefficients
  
  SUBROUTINE tridiagonal_solver(a,b,c,d,w,n)
    ! Tridiagnal solver, using Thomas' algorithm.
    INTEGER :: i,n
    REAL(wp) :: a(0:n),b(0:n),c(0:n),d(0:n),w(0:n)
    ! First eliminate all a_i's
    DO i=1,n
       ! a(i) = a(i) - b(i-1)*a(i)/b(i-1) 
       b(i) = b(i) - c(i-1)*a(i)/b(i-1) ! should really chack here that b(i-1)\=0
       ! c(i) = c(i) - 0*a(i)/b(i-1)
       d(i) = d(i) - d(i-1)*a(i)/b(i-1)
    END DO
    ! now back substitute
    w(n) = d(n)/b(n)
    DO i=n-1,0,-1
       ! each equation is of the form
       !     b_i w_i + c_i w_i+1 = d_i
       w(i) = (d(i) - c(i)*w(i+1))/b(i)
    ENDDO
       
  END SUBROUTINE tridiagonal_solver

END MODULE Finite_difference

PROGRAM fd_solver
  ! we need all the variables from global data
  USE global_data, ONLY : wp=>dp,kappa,h,x_min,x_max
  ! RungeKutta_method gives us our rungekutta method function
  USE finite_difference
  IMPLICIT NONE
  
  INTEGER :: no_of_points,i
  REAL(wp), ALLOCATABLE :: a(:),b(:),c(:),d(:),w(:),x(:)

  
  ! Setup initial conditions
  ! set x_0 = 0
  x_min = 0._wp
  ! set x_n = 1
  x_max = 1._wp

  ! Set the number of steps
  no_of_points = 1000
  ! Set the parameter kappa from the ODE
  kappa = 5._wp
  ! Set the step size
  h = (x_max-x_min)/dble(no_of_points)

  ! allocate arrays 
  ALLOCATE(a(0:no_of_points),b(0:no_of_points),c(0:no_of_points),d(0:no_of_points),w(0:no_of_points),x(0:no_of_points))
  ! setup the array x
  DO i=0,no_of_points
     x(i) = x_min + dble(i)*h
  END DO

  ! Output some info to screen
  WRITE(6,*)' Run finite difference code with ::',no_of_points,' points.'
  
  PRINT*,' Setup LHS of the matrix equation.'
  CALL  setup_lhs_coefficients(a,b,c,x,no_of_points)
  
  PRINT*,' Setup RHS of matrix equation.'
  CALL  setup_rhs_coefficients(d,x,no_of_points)

  PRINT*,' Solve using Thomas algorithm.'
  CALL  tridiagonal_solver(a,b,c,d,w,no_of_points)

  PRINT*,' Print results to file.'
  CALL print_solution(x,w,no_of_points,"results.dat",10,"(2(1X,E15.8))")
  
END PROGRAM fd_solver

SUBROUTINE print_solution(x,y,n,filename,unit,format_str)
  USE Global_data, ONLY : wp=>dp
  ! Subroutine to print x versus y to file

  INTEGER :: i,n,unit
  REAL(wp) :: x(0:n),y(0:n)
  CHARACTER(LEN=*) :: filename,format_str

  OPEN(UNIT=unit,FILE=filename)

  DO i=0,n
     WRITE(10,format_str)x(i),y(i)
  END DO

  CLOSE(UNIT=unit)

END SUBROUTINE print_solution