program main
!=======================================================================
! Chapter 10, code 2: a linear BVP with different boundary conditions.
! Solve y'' + p(x)*y' + q(x)*y = r(x) on a uniform grid.
! Method: centered finite differences and the Thomas algorithm.
! Derivative boundaries use ghost-point elimination, not a first-order
! reconstruction of the endpoint value. All endpoint values are stored.
!
! Example: dimensionless steady heat conduction in a rod,
!   T'' = -sin(pi*x), 0 <= x <= 1.
! The source represents sinusoidally varying internal heating.
! Default: T(0) = 0, T'(1) = 0 (an insulated right end).
! y denotes T in this example. Exact solutions are supplied for all keys.
!
! To solve another problem, edit the inputs and coefficients below;
! replace exact_solution or set check_exact = .false.
!=======================================================================
implicit none
integer, parameter :: n = 20, data_unit = 20
integer key, i
double precision a, b, left_value, right_value, robin_alpha, robin_beta
double precision x(0:n), y(0:n), h, exact, error, max_error
double precision exact_solution
logical check_exact
external coefficients, exact_solution

!-------------------- Problem and grid: edit here -----------------------
a = 0.0d0
b = 1.0d0
key = 2
left_value = 0.0d0
right_value = 0.0d0
robin_alpha = 1.0d0
robin_beta = 1.0d0
check_exact = .true.
! Change n in the declaration above to refine the grid.
!
! Meaning of key and the two boundary data:
!   1: y(a)  = left_value; y(b)  = right_value
!   2: y(a)  = left_value; y'(b) = right_value
!   3: y'(a) = left_value; y(b)  = right_value
!   4: y(a)  = left_value;
!      robin_alpha*y(b) + robin_beta*y'(b) = right_value
!
! Derivatives always mean dy/dx, not the outward normal derivative.
! For a rod's convecting right end, key 4 expresses
!   k*T'(b) + H*T(b) = H*T_environment,
! where k is conductivity and H is a heat-transfer coefficient.
! Use their dimensionless counterparts in this normalized example.
!-----------------------------------------------------------------------
if (n < 2 .or. b <= a) stop 'Use n >= 2 and b > a.'
if (key < 1 .or. key > 4) stop 'Choose a boundary key from 1 to 4.'
if (key == 4) then
    if (robin_beta == 0.0d0) stop 'Use key 1 for a fixed right-end value.'
    if (check_exact) then
        if (abs(robin_alpha*(b-a)+robin_beta) <= epsilon(1.0d0)* &
            max(abs(robin_alpha*(b-a)),abs(robin_beta))) &
            stop 'The supplied example has no unique solution for these Robin data.'
    end if
end if

h = (b-a)/dble(n)
do i = 0,n
    x(i) = a+dble(i)*h
end do
x(n) = b
call finite_difference(coefficients,n,x,key,left_value,right_value, &
                       robin_alpha,robin_beta,y)

open(unit=data_unit,file='ch10_boundary_conditions.dat', &
     status='replace',action='write')
write(*,'(a)') 'Linear boundary value problem: boundary-condition choices'
write(*,'(a,i8)') 'Boundary key:  ', key
write(*,'(a,i8)') 'Grid intervals:', n
write(*,'(a,f12.6)') 'Grid spacing:  ', h
select case (key)
case (1)
    write(*,'(a)') 'Fixed values at both ends.'
case (2)
    write(*,'(a)') 'Fixed left value; prescribed right derivative.'
case (3)
    write(*,'(a)') 'Prescribed left derivative; fixed right value.'
case (4)
    write(*,'(a)') 'Fixed left value; right Robin condition.'
end select
write(*,'(a)') ''
if (check_exact) then
    write(*,'(a10,2(1x,a16),1x,a12)') 'x','numerical y','exact y','abs error'
    write(data_unit,'(a)') '# x numerical_y exact_y abs_error'
else
    write(*,'(a10,1x,a16)') 'x','numerical y'
    write(data_unit,'(a)') '# x numerical_y'
end if
max_error = 0.0d0
do i = 0,n
    if (check_exact) then
        exact = exact_solution(x(i),a,b,key,left_value,right_value,robin_alpha,robin_beta)
        error = abs(y(i)-exact)
        max_error = max(max_error,error)
        write(*,'(f10.6,2(1x,es16.8),1x,es12.4)') x(i),y(i),exact,error
        write(data_unit,'(4(1x,es24.16))') x(i),y(i),exact,error
    else
        write(*,'(f10.6,1x,es16.8)') x(i),y(i)
        write(data_unit,'(2(1x,es24.16))') x(i),y(i)
    end if
end do
close(data_unit)
write(*,'(a)') ''
if (check_exact) write(*,'(a,es12.4)') 'Largest grid error: ', max_error
write(*,'(a)') 'Data file: ch10_boundary_conditions.dat'
end program main


subroutine coefficients(x,p,q,r)
!=======================================================================
! Supply y'' + p(x)*y' + q(x)*y = r(x).
! Here y is dimensionless temperature in a uniformly conducting rod.
! For another model, change these three coefficients.
!=======================================================================
implicit none
double precision x, p, q, r, pi
pi = acos(-1.0d0)
p = 0.0d0
q = 0.0d0
r = -sin(pi*x)
end subroutine coefficients


double precision function exact_solution(x,a,b,key,left_value,right_value,alpha,beta)
!=======================================================================
! For y'' = -sin(pi*x), y = sin(pi*x)/pi**2 + c*(x-a) + d.
! Determine c and d from the boundary conditions selected by key.
! Default key 2 on [0,1]: y = sin(pi*x)/pi**2 + x/pi.
!=======================================================================
implicit none
integer key
double precision x, a, b, left_value, right_value, alpha, beta
double precision pi, pa, pb, c, d
pi = acos(-1.0d0)
pa = sin(pi*a)/pi**2
pb = sin(pi*b)/pi**2
select case (key)
case (1)
    d = left_value-pa
    c = (right_value-pb-d)/(b-a)
case (2)
    d = left_value-pa
    c = right_value-cos(pi*b)/pi
case (3)
    c = left_value-cos(pi*a)/pi
    d = right_value-pb-c*(b-a)
case (4)
    d = left_value-pa
    c = (right_value-alpha*(pb+d)-beta*cos(pi*b)/pi)/(alpha*(b-a)+beta)
end select
exact_solution = sin(pi*x)/pi**2+c*(x-a)+d
end function exact_solution


subroutine finite_difference(fcn,n,x,key,left_value,right_value,alpha,beta,y)
!=======================================================================
! Interior equations are identical to Code 1.
! Here all n+1 values, including endpoints, are unknowns in the system.
! A fixed endpoint gets a row containing just y = prescribed value.
!
! For a derivative boundary introduce one ghost point outside the grid.
! At the left end:  y(-1) = y(1)-2*h*y'(a).
! At the right end: y(n+1) = y(n-1)+2*h*y'(b).
! Substitute into the centered differential equation at the endpoint.
! This removes the ghost point and keeps the system tridiagonal.
!
! For the right Robin condition use y'(b) = (right_value-alpha*y(n))/beta
! in the same right-end equation. Only that boundary row changes.
!=======================================================================
implicit none
integer n, key, i
double precision x(0:n), y(0:n), left_value, right_value, alpha, beta
double precision lower(0:n), diagonal(0:n), upper(0:n), rhs(0:n)
double precision h, p, q, r, factor
external fcn
h = (x(n)-x(0))/dble(n)
lower = 0.0d0
upper = 0.0d0

do i = 1,n-1
    call fcn(x(i),p,q,r)
    lower(i) = 1.0d0-h*p/2.0d0
    diagonal(i) = -2.0d0+h*h*q
    upper(i) = 1.0d0+h*p/2.0d0
    rhs(i) = h*h*r
end do

! Left boundary: derivative for key 3; fixed value otherwise.
if (key == 3) then
    call fcn(x(0),p,q,r)
    diagonal(0) = -2.0d0+h*h*q
    upper(0) = 2.0d0
    rhs(0) = h*h*r+(2.0d0*h-h*h*p)*left_value
else
    diagonal(0) = 1.0d0
    rhs(0) = left_value
end if

! Right boundary: fixed value, derivative, or Robin condition.
select case (key)
case (1,3)
    diagonal(n) = 1.0d0
    rhs(n) = right_value
case (2)
    call fcn(x(n),p,q,r)
    lower(n) = 2.0d0
    diagonal(n) = -2.0d0+h*h*q
    rhs(n) = h*h*r-(2.0d0*h+h*h*p)*right_value
case (4)
    call fcn(x(n),p,q,r)
    factor = (2.0d0*h+h*h*p)/beta
    lower(n) = 2.0d0
    diagonal(n) = -2.0d0+h*h*q-factor*alpha
    rhs(n) = h*h*r-factor*right_value
end select
call thomas(n+1,lower,diagonal,upper,rhs,y)
end subroutine finite_difference


subroutine thomas(n,lower,diagonal,upper,rhs,y)
!=======================================================================
! Thomas algorithm, as in Code 1. diagonal and rhs are overwritten.
! Arrays passed with indices 0:n in the caller are indexed 1:n here:
! explicit-shape dummy arrays refer to the same consecutive elements.
! No pivoting: a zero or near-zero pivot stops the calculation.
!=======================================================================
implicit none
integer n, i
double precision lower(n), diagonal(n), upper(n), rhs(n), y(n)
double precision multiplier, scale
do i = 1,n
    scale = max(abs(lower(i)),abs(diagonal(i)),abs(upper(i)))
    if (abs(diagonal(i)) <= epsilon(1.0d0)*scale) &
        stop 'Thomas algorithm: zero or near-zero pivot.'
    if (i == n) exit
    multiplier = lower(i+1)/diagonal(i)
    diagonal(i+1) = diagonal(i+1)-multiplier*upper(i)
    rhs(i+1) = rhs(i+1)-multiplier*rhs(i)
end do
y(n) = rhs(n)/diagonal(n)
do i = n-1,1,-1
    y(i) = (rhs(i)-upper(i)*y(i+1))/diagonal(i)
end do
end subroutine thomas