program main
!=======================================================================
! Chapter 10, code 0: shooting for a second-order boundary value problem.
! Solve y'' = f(x,y,y'), with y(a) = alpha and y(b) = beta.
! An unknown initial slope is found by secant iteration. Each trial is
! an initial value problem, integrated with adaptive Fehlberg RKF45.
! Based on the shooting structure in shoot2.f90 by Alexander Godunov.
!
! Example: y'' = y, y(0) = y(1) = 1.
! Exact solution: cosh(x-0.5)/cosh(0.5).
! To solve another problem:
!   1. Edit the interval, boundary values, and two slope guesses below.
!   2. Edit f(x,y,dy) immediately below the driver.
!   3. Disable check_exact or replace exact_solution.
!=======================================================================
implicit none
integer, parameter :: capacity = 100001, data_unit = 20
! Allocate storage on the heap, rather than using large stack arrays.
double precision, allocatable :: x(:), y(:,:)
double precision a, b, alpha, beta, slope1, slope2, slope
double precision h_first, h_max, ivp_tolerance, boundary_tolerance
double precision residual, exact, error, max_error
double precision f, exact_solution
integer i, trials, accepted, rejected
logical check_exact
external f, exact_solution

!-------------------- Problem and controls: edit here --------------------
a = 0.0d0
b = 1.0d0
alpha = 1.0d0                   ! y(a)
beta = 1.0d0                    ! y(b)
slope1 = 0.0d0                  ! First guess for y'(a)
slope2 = 1.0d0                  ! Second guess for y'(a)
h_first = 0.05d0
h_max = 0.05d0                  ! Maximum spacing between output points
ivp_tolerance = 1.0d-10         ! Absolute LOCAL error tolerance for RKF45
boundary_tolerance = 1.0d-8     ! Required |y(b)-beta|
check_exact = .true.
!------------------------------------------------------------------------
allocate(x(capacity),y(2,capacity))
write(*,'(a)') 'Boundary value problem: shooting with secant iteration and RKF45'
write(*,'(a,es12.4)') 'RKF45 tolerance:   ', ivp_tolerance
write(*,'(a,es12.4)') 'Boundary tolerance:', boundary_tolerance
call shoot2(f,a,b,alpha,beta,slope1,slope2,h_first,h_max, &
            ivp_tolerance,boundary_tolerance,capacity,x,y,slope,residual, &
            trials,accepted,rejected)

! Only the converged trajectory is printed and written to the file.
open(unit=data_unit,file='ch10_shooting.dat',status='replace',action='write')
max_error = 0.0d0
write(*,'(a)') ''
if (check_exact) then
    write(*,'(a10,3(1x,a16),1x,a12)') 'x','numerical y','dy/dx','exact y','abs error'
    write(data_unit,'(a)') '# x y dy_dx exact_y abs_error'
else
    write(*,'(a10,2(1x,a16))') 'x','numerical y','dy/dx'
    write(data_unit,'(a)') '# x y dy_dx'
end if
do i = 1,accepted+1
    if (check_exact) then
        exact = exact_solution(x(i))
        error = abs(y(1,i)-exact)
        max_error = max(max_error,error)
        write(*,'(f10.6,3(1x,es16.8),1x,es12.4)') x(i),y(:,i),exact,error
        write(data_unit,'(5(1x,es24.16))') x(i),y(:,i),exact,error
    else
        write(*,'(f10.6,2(1x,es16.8))') x(i),y(:,i)
        write(data_unit,'(3(1x,es24.16))') x(i),y(:,i)
    end if
end do
close(data_unit)
write(*,'(a)') ''
write(*,'(a,i8)') 'Shooting trials:      ', trials
write(*,'(a,f16.10)') 'Initial slope:        ', slope
write(*,'(a,es16.6)') 'Boundary residual:    ', residual
write(*,'(a,i8)') 'Accepted steps:       ', accepted
write(*,'(a,i8)') 'Rejected trials:      ', rejected
if (check_exact) write(*,'(a,es16.6)') 'Largest sampled error:', max_error
write(*,'(a)') 'Step counts refer to the final shooting trial.'
write(*,'(a)') 'Data file: ch10_shooting.dat'
end program main

double precision function f(x,y,dy)
! The original second-order equation: y'' = f(x,y,y').
! x and dy are available for equations that depend on them.
implicit none
double precision x, y, dy
f = y
end function f

double precision function exact_solution(x)
! Analytic solution for the supplied example and boundary values.
implicit none
double precision x
exact_solution = cosh(x-0.5d0)/cosh(0.5d0)
end function exact_solution

subroutine shoot2(f,a,b,alpha,beta,slope1,slope2,h_first,h_max, &
                  ivp_tolerance,boundary_tolerance,capacity,grid,solution, &
                  slope,residual,trials,accepted,rejected)
!=======================================================================
! Treat s = y'(a) as an unknown. For each trial integrate from a to b,
! starting with y(a) = alpha and y'(a) = s.
! The mismatch R(s) = y(b;s)-beta must vanish.
! Two successive slopes and residuals determine the next secant slope.
! The final integration leaves the converged trajectory in the arrays.
! Secant iteration is not bracketed; convergence depends on the guesses.
! RKF45 controls local integration error. A small boundary residual alone
! does not establish accuracy of the whole solution: tighten the RKF45
! tolerance and compare trajectories when checking a new problem.
!=======================================================================
implicit none
integer capacity, trials, accepted, rejected
integer, parameter :: max_trials = 30
double precision f, a, b, alpha, beta, slope1, slope2, h_first, h_max
double precision ivp_tolerance, boundary_tolerance
double precision grid(capacity), solution(2,capacity), slope, residual
double precision previous_slope, previous_residual, next_slope, denominator
double precision initial(2), final(2)
external f
if (boundary_tolerance <= 0.0d0 .or. slope1 == slope2) &
    stop 'Use positive boundary tolerance and two distinct slope guesses.'
write(*,'(a8,2(1x,a18))') 'trial','initial slope','boundary residual'
slope = slope1
previous_slope = slope1
previous_residual = 0.0d0

do trials = 1,max_trials
    initial = (/ alpha, slope /)
    call rkf45_system(f,a,initial,b,h_first,h_max,ivp_tolerance, &
                      capacity,grid,solution,final,accepted,rejected)
    residual = final(1)-beta
    write(*,'(i8,2(1x,es18.10))') trials,slope,residual
    if (abs(residual) <= boundary_tolerance) return
    if (trials == max_trials) exit
    if (trials == 1) then
        next_slope = slope2
    else
        denominator = residual-previous_residual
        if (abs(denominator) <= epsilon(denominator)* &
            max(abs(residual),abs(previous_residual),tiny(denominator))) &
            stop 'Secant residuals are indistinguishable; change the guesses.'
        next_slope = slope-residual*(slope-previous_slope)/denominator
        if (.not. (abs(next_slope) <= huge(next_slope))) stop 'Nonfinite secant slope.'
        if (next_slope == slope) stop 'Secant slope cannot advance; check tolerances.'
    end if
    previous_slope = slope
    previous_residual = residual
    slope = next_slope
end do
stop 'Shooting did not converge within the trial limit.'
end subroutine shoot2

subroutine equations(f,x,y,dy)
! Write the second-order equation as two first-order equations:
! y(1) = y, y(2) = y'; therefore y(1)' = y(2), y(2)' = f(x,y,y').
implicit none
double precision f, x, y(2), dy(2)
external f
dy(1) = y(2)
dy(2) = f(x,y(1),y(2))
! HUGE gives the largest finite value of this floating-point type.
! Avoid IEEE_ARITHMETIC, which causes a linker conflict in Absoft 21.
if (.not. all(abs(dy) <= huge(1.0d0))) stop 'Nonfinite derivative in trial integration.'
end subroutine equations

subroutine rkf45_system(fcn,t0,x0,t_end,h_first,h_max,tolerance, &
                        capacity,grid,solution,xf,accepted,rejected)
!=======================================================================
! Adaptive Fehlberg RKF45 for the two first-order equations.
! Retain initial and accepted points. Trial trajectories produce no output.
! Advance with the fifth-order estimate; use the largest absolute embedded
! component error for acceptance. All adaptive step decisions are here.
! One absolute tolerance assumes components of comparable magnitudes.
!=======================================================================
implicit none
integer, parameter :: n = 2
integer capacity, accepted, rejected, attempts, j
double precision t0, x0(n), t_end, h_first, h_max, tolerance, xf(n)
double precision fcn, grid(capacity), solution(n,capacity)
double precision t, tf, h, err, component_error, factor
double precision x(n), work(n), dx(n), trial(n)
double precision k1(n), k2(n), k3(n), k4(n), k5(n), k6(n)
integer, parameter :: max_attempts = 100000
external fcn

if (capacity < 2 .or. t_end <= t0 .or. h_first <= 0.0d0 .or. h_max <= 0.0d0 .or. tolerance <= 0.0d0) &
    stop 'Use capacity >= 2, b > a, positive steps, and positive tolerance.'
t = t0
h = h_first
x = x0
accepted = 0
rejected = 0
attempts = 0
grid(1) = t
solution(:,1) = x

do while (t < t_end)
    if (attempts >= max_attempts) stop 'Maximum number of trials exceeded.'
    h = min(h,h_max,t_end-t)
    tf = t+h
    if (tf <= t) stop 'Step too small to advance time.'
    attempts = attempts+1

    ! Stage 1, then prepare the state for stage 2. Every k includes h.
    call equations(fcn,t,x,dx)
    do j = 1,n
        k1(j) = h*dx(j)
        work(j) = x(j)+k1(j)/4.0d0
    end do
    ! Stage 2: all equations use the SAME intermediate state work.
    call equations(fcn,t+h/4.0d0,work,dx)
    do j = 1,n
        k2(j) = h*dx(j)
        work(j) = x(j)+3.0d0*k1(j)/32.0d0+9.0d0*k2(j)/32.0d0
    end do
    call equations(fcn,t+3.0d0*h/8.0d0,work,dx)
    do j = 1,n
        k3(j) = h*dx(j)
        work(j) = x(j)+1932.0d0*k1(j)/2197.0d0 &
                  -7200.0d0*k2(j)/2197.0d0+7296.0d0*k3(j)/2197.0d0
    end do
    call equations(fcn,t+12.0d0*h/13.0d0,work,dx)
    do j = 1,n
        k4(j) = h*dx(j)
        work(j) = x(j)+439.0d0*k1(j)/216.0d0-8.0d0*k2(j) &
                  +3680.0d0*k3(j)/513.0d0-845.0d0*k4(j)/4104.0d0
    end do
    call equations(fcn,t+h,work,dx)
    do j = 1,n
        k5(j) = h*dx(j)
        work(j) = x(j)-8.0d0*k1(j)/27.0d0+2.0d0*k2(j) &
                  -3544.0d0*k3(j)/2565.0d0+1859.0d0*k4(j)/4104.0d0 &
                  -11.0d0*k5(j)/40.0d0
    end do
    call equations(fcn,t+h/2.0d0,work,dx)
    do j = 1,n
        k6(j) = h*dx(j)
    end do

    ! Fifth-order candidate and largest embedded component error.
    err = 0.0d0
    do j = 1,n
        trial(j) = x(j)+16.0d0*k1(j)/135.0d0+6656.0d0*k3(j)/12825.0d0 &
                   +28561.0d0*k4(j)/56430.0d0-9.0d0*k5(j)/50.0d0+2.0d0*k6(j)/55.0d0
        component_error = abs(k1(j)/360.0d0-128.0d0*k3(j)/4275.0d0 &
                         -2197.0d0*k4(j)/75240.0d0+k5(j)/50.0d0+2.0d0*k6(j)/55.0d0)
        err = max(err,component_error)
    end do
    if (.not. all(abs(trial) <= huge(1.0d0)) .or. .not. (abs(err) <= huge(err))) &
        stop 'Nonfinite RKF45 result: check the model and slope guesses.'
    if (err <= tolerance) then
        if (h == t_end-t) tf = t_end
        t = tf
        x = trial
        accepted = accepted+1
        if (accepted+1 > capacity) stop 'Increase trajectory storage capacity.'
        grid(accepted+1) = t
        solution(:,accepted+1) = x
    else
        rejected = rejected+1  ! Keep both t and x unchanged.
    end if
    if (err == 0.0d0) then
        factor = 2.0d0
    else
        factor = 0.9d0*(tolerance/err)**0.2d0
    end if
    if (err <= tolerance) then
        h = h*min(2.0d0,factor)
    else
        h = h*max(0.2d0,factor)
    end if
end do
xf = x
end subroutine rkf45_system