program main
!=======================================================================
! Chapter 9, code 3: a second-order ODE, solved with adaptive RKF45.
! Damped spring: mass*x'' + damping*x' + spring*x = 0.
! Introduce velocity v = x':
!   x' = v
!   v' = -(spring/mass)*x - (damping/mass)*v
! Store position in y(1), velocity in y(2).
!
! To use another physical model:
!   1. Edit the parameters and initial data below.
!   2. Edit the acceleration equation in equations.
!   3. Disable check_exact or replace the exact-solution calculation.
! The adaptive RKF45 stages and step control need no changes.
!=======================================================================
implicit none
integer, parameter :: n = 2, data_unit = 20
double precision mass, damping, spring
double precision t0, t_end, h_first, h_max, tolerance
double precision y0(n), yf(n), exact(n)
double precision gamma, omega, a, b, tau, decay
integer accepted, rejected
logical check_exact
external equations

!-------------------- Parameters and initial data: edit here -------------
mass = 1.0d0
spring = 1.0d0
damping = 0.2d0
t0 = 0.0d0
t_end = 20.0d0
y0(1) = 1.0d0           ! Initial position
y0(2) = 0.0d0           ! Initial velocity
h_first = 0.1d0
h_max = 0.1d0          ! Maximum step: enough output points per oscillation
tolerance = 1.0d-8      ! Absolute local tolerance for each component
check_exact = .true.    ! Analytic comparison for an underdamped spring
!------------------------------------------------------------------------
if (mass <= 0.0d0 .or. spring <= 0.0d0 .or. damping < 0.0d0) &
    stop 'Use positive mass and spring, and nonnegative damping.'
if (check_exact .and. damping**2 >= 4.0d0*mass*spring) &
    stop 'Exact comparison requires an underdamped spring; disable check_exact.'

open(unit=data_unit, file='ch09_rkf45_spring.dat', &
     status='replace', action='write')
write(*,'(a)') 'Adaptive RKF45: damped spring'
write(*,'(a,es12.4)') 'Tolerance: ', tolerance
call rkf45_system(equations,n,t0,y0,t_end,h_first,h_max,tolerance,mass,damping,spring,data_unit, &
                  yf,accepted,rejected)
close(data_unit)
write(*,'(a,i8)') 'Accepted steps: ', accepted
write(*,'(a,i8)') 'Rejected trials:', rejected

if (check_exact) then
    ! Exact underdamped solution, including zero damping.
    ! tau measures time from the initial point.
    ! x = exp(-gamma*tau)*(a*cos(omega*tau)+b*sin(omega*tau))
    gamma = damping/(2.0d0*mass)
    omega = sqrt(spring/mass-gamma**2)
    a = y0(1)
    b = (y0(2)+gamma*a)/omega
    tau = t_end-t0
    decay = exp(-gamma*tau)
    exact(1) = decay*(a*cos(omega*tau)+b*sin(omega*tau))
    exact(2) = -gamma*exact(1) &
               +decay*omega*(-a*sin(omega*tau)+b*cos(omega*tau))
    write(*,'(a10,2(1x,a18),1x,a12)') 'variable','numerical','exact','abs error'
    write(*,'(a10,2(1x,es18.10),1x,es12.4)') &
        'position',yf(1),exact(1),abs(yf(1)-exact(1))
    write(*,'(a10,2(1x,es18.10),1x,es12.4)') &
        'velocity',yf(2),exact(2),abs(yf(2)-exact(2))
end if
write(*,'(a)') 'Data file: ch09_rkf45_spring.dat'

end program main

subroutine equations(t,y,dy,n,mass,damping,spring)
!=======================================================================
! y(1) = position, y(2) = velocity.
! The driver passes mass, damping, and spring explicitly.
! For another force, change the acceleration dy(2).
! t is available for a force that depends explicitly on time.
!=======================================================================
implicit none
integer n
double precision t, y(n), dy(n), mass, damping, spring
dy(1) = y(2)
dy(2) = -(spring/mass)*y(1) - (damping/mass)*y(2)
end subroutine equations

subroutine rkf45_system(fcn,n,t0,x0,t_end,h_first,h_max,tolerance,mass,damping,spring,data_unit, &
                        xf,accepted,rejected)
!=======================================================================
! Solve the two position/velocity equations with adaptive Fehlberg RKF45.
! All step decisions are inside this subroutine.
! h_max caps the output spacing; the tolerance can require smaller steps.
! fcn(t,x,dx,n,mass,damping,spring) returns the derivatives in dx.
! Input: initial data, interval, step, tolerance, physical parameters,
!        and opened data_unit.
! Output: xf(n), accepted and rejected counts.
! Screen and file output include the initial point and accepted points.
!
! Advance with the fifth-order estimate. The largest absolute embedded
! component error controls acceptance, using ONE absolute tolerance.
! This is suitable for components of comparable scales. Variables with
! different units or magnitudes may require individual scaled tolerances.
! The tolerance controls local error, not accumulated global error.
!=======================================================================
implicit none
integer n, data_unit, accepted, rejected, attempts, j
double precision t0, x0(n), t_end, h_first, h_max, tolerance, xf(n)
double precision mass, damping, spring
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
double precision, parameter :: h_min = 1.0d-12
external fcn

if (n < 1 .or. t_end <= t0 .or. h_first <= 0.0d0 .or. h_max <= 0.0d0 .or. tolerance <= 0.0d0) &
    stop 'Use n > 0, t_end > t0, positive steps, and positive tolerance.'
! Output columns: time, position, velocity, and the accepted step size.
write(*,'(a10,3(1x,a18))') 't','position','velocity','h'
write(data_unit,'(a)') '# t position velocity h'

t = t0
h = h_first
x = x0
accepted = 0
rejected = 0
attempts = 0
write(*,'(f10.6,3(1x,es18.10))') t,x,0.0d0
write(data_unit,'(4(1x,es24.16))') t,x,0.0d0

do while (t < t_end)
    if (t_end-t <= 64.0d0*epsilon(t)*max(1.0d0,abs(t_end))) exit
    if (attempts >= max_attempts) stop 'Maximum number of trials exceeded.'
    h = min(h,h_max,t_end-t)
    tf = t+h
    if (h < h_min .or. 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 fcn(t,x,dx,n,mass,damping,spring)
    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 fcn(t+h/4.0d0,work,dx,n,mass,damping,spring)
    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 fcn(t+3.0d0*h/8.0d0,work,dx,n,mass,damping,spring)
    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 fcn(t+12.0d0*h/13.0d0,work,dx,n,mass,damping,spring)
    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 fcn(t+h,work,dx,n,mass,damping,spring)
    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 fcn(t+h/2.0d0,work,dx,n,mass,damping,spring)
    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 (err <= tolerance) then
        if (h == t_end-t) tf = t_end
        t = tf
        x = trial
        accepted = accepted+1
        write(*,'(f10.6,3(1x,es18.10))') t,x,h
        write(data_unit,'(4(1x,es24.16))') t,x,h
    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